witbitz-code 1.6.9__tar.gz → 1.6.10__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (102) hide show
  1. {witbitz_code-1.6.9/src/witbitz_code.egg-info → witbitz_code-1.6.10}/PKG-INFO +2 -1
  2. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/pyproject.toml +3 -1
  3. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/__init__.py +1 -1
  4. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/confidential.py +7 -1
  5. witbitz_code-1.6.10/src/witbitz_code/images.py +125 -0
  6. {witbitz_code-1.6.9 → witbitz_code-1.6.10/src/witbitz_code.egg-info}/PKG-INFO +2 -1
  7. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code.egg-info/SOURCES.txt +2 -0
  8. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code.egg-info/requires.txt +1 -0
  9. witbitz_code-1.6.10/tests/test_images.py +102 -0
  10. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/LICENSE +0 -0
  11. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/README.md +0 -0
  12. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/setup.cfg +0 -0
  13. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/__main__.py +0 -0
  14. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/_js.py +0 -0
  15. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/account.py +0 -0
  16. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/asks.py +0 -0
  17. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/attachments.py +0 -0
  18. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/auto.py +0 -0
  19. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/auto_runner.py +0 -0
  20. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/cli.py +0 -0
  21. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/codekeys.py +0 -0
  22. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/codetools.py +0 -0
  23. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/compress.py +0 -0
  24. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/computers.py +0 -0
  25. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/confidential_models.json +0 -0
  26. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/connector.py +0 -0
  27. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/devicelink.py +0 -0
  28. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/dirs.py +0 -0
  29. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/folders.py +0 -0
  30. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/gateway.py +0 -0
  31. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/icons/witbitz-code-wsl.ico +0 -0
  32. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/icons/witbitz-code.ico +0 -0
  33. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/keys.py +0 -0
  34. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/link.py +0 -0
  35. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/mkdir.py +0 -0
  36. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/models.py +0 -0
  37. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/net.py +0 -0
  38. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/offered_models.json +0 -0
  39. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/opencode_plugins/notes-init.md +0 -0
  40. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/opencode_plugins/witbitz-media-server.mjs +0 -0
  41. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/opencode_plugins/witbitz-notes.js +0 -0
  42. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/opencode_plugins/witbitz-progress.js +0 -0
  43. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/outputs.py +0 -0
  44. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/pair.py +0 -0
  45. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/pairings.py +0 -0
  46. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/payers.py +0 -0
  47. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/plugins.py +0 -0
  48. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/policy.py +0 -0
  49. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/ports.py +0 -0
  50. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/qr.py +0 -0
  51. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/receipt.py +0 -0
  52. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/relay.py +0 -0
  53. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/seen.py +0 -0
  54. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/tinfoil_read.py +0 -0
  55. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/tools_catalog.json +0 -0
  56. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/tools_probe.py +0 -0
  57. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/toolsetup.py +0 -0
  58. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/ui.py +0 -0
  59. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/ui_static/app.css +0 -0
  60. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/ui_static/app.js +0 -0
  61. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/ui_static/index.html +0 -0
  62. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/uitools.py +0 -0
  63. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/uninstall.py +0 -0
  64. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/winagent.py +0 -0
  65. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/winbg.py +0 -0
  66. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/winconsole.py +0 -0
  67. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/wizard.py +0 -0
  68. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code/wslstart.py +0 -0
  69. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code.egg-info/dependency_links.txt +0 -0
  70. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code.egg-info/entry_points.txt +0 -0
  71. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/src/witbitz_code.egg-info/top_level.txt +0 -0
  72. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_asks.py +0 -0
  73. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_attachments.py +0 -0
  74. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_auto.py +0 -0
  75. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_cli.py +0 -0
  76. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_codekeys.py +0 -0
  77. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_codetools.py +0 -0
  78. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_confidential.py +0 -0
  79. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_connector_e2e.py +0 -0
  80. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_devicelink.py +0 -0
  81. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_dirs.py +0 -0
  82. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_files_and_compress.py +0 -0
  83. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_folders.py +0 -0
  84. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_outputs.py +0 -0
  85. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_pair_flows.py +0 -0
  86. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_plugins.py +0 -0
  87. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_policy.py +0 -0
  88. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_ports.py +0 -0
  89. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_ports_win.py +0 -0
  90. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_registry.py +0 -0
  91. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_relay_codec.py +0 -0
  92. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_setup.py +0 -0
  93. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_tools_probe.py +0 -0
  94. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_toolsetup.py +0 -0
  95. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_ui.py +0 -0
  96. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_uitools.py +0 -0
  97. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_uninstall.py +0 -0
  98. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_winagent.py +0 -0
  99. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_winbg.py +0 -0
  100. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_winconsole.py +0 -0
  101. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_windows_files.py +0 -0
  102. {witbitz_code-1.6.9 → witbitz_code-1.6.10}/tests/test_wslstart.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: witbitz-code
3
- Version: 1.6.9
3
+ Version: 1.6.10
4
4
  Summary: Reach OpenCode on this computer from the Witbitz Spaces Code section, through an end-to-end sealed relay.
5
5
  Author: Witbitz
6
6
  License-Expression: MIT
@@ -19,6 +19,7 @@ Requires-Dist: websockets>=13
19
19
  Requires-Dist: httpx>=0.25
20
20
  Requires-Dist: segno>=1.5
21
21
  Requires-Dist: tinfoil<0.15,>=0.14
22
+ Requires-Dist: pillow>=10
22
23
  Provides-Extra: test
23
24
  Requires-Dist: pytest>=8; extra == "test"
24
25
  Requires-Dist: pyflakes>=3; extra == "test"
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "witbitz-code"
7
- version = "1.6.9"
7
+ version = "1.6.10"
8
8
  description = "Reach OpenCode on this computer from the Witbitz Spaces Code section, through an end-to-end sealed relay."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -26,6 +26,8 @@ dependencies = [
26
26
  # Tinfoil's own verifier (attestation + Sigstore + TLS pinning) — the confidential proxy reads images and files a
27
27
  # text-only confidential model cannot, in Tinfoil's attested enclave (tinfoil_read.py). Pinned to a series on purpose.
28
28
  "tinfoil>=0.14,<0.15",
29
+ # Older screenshots are re-encoded smaller before a request leaves (images.py): TrustedRouter drops a ~10 MB request.
30
+ "pillow>=10",
29
31
  ]
30
32
 
31
33
  [project.optional-dependencies]
@@ -4,4 +4,4 @@ The Python implementation of tools/witbitz-code.mjs (docs/opencode-relay.md phas
4
4
  vectors and cross-implementation tests.
5
5
  """
6
6
 
7
- __version__ = "1.6.9"
7
+ __version__ = "1.6.10"
@@ -26,7 +26,7 @@ from typing import Any, Callable
26
26
 
27
27
  import httpx
28
28
 
29
- from . import keys, models as mdl, receipt as rcp, tinfoil_read
29
+ from . import images, keys, models as mdl, receipt as rcp, tinfoil_read
30
30
  from .gateway import verify_gateway_attestation
31
31
 
32
32
  TR_UPSTREAM = "https://api.trustedrouter.com"
@@ -577,6 +577,12 @@ class ConfidentialProxy:
577
577
  return
578
578
  req.headers = [(k, v) for k, v in req.headers if k.lower() not in ("authorization", "x-api-key")] + [("authorization", "Bearer " + cred["trustedrouter"])]
579
579
  tinfoil = cred.get("tinfoil") or ""
580
+ if isinstance(body, dict):
581
+ shrunk, rep = await asyncio.to_thread(images.shrink_older_images, body)
582
+ if shrunk is not body:
583
+ body, req.body = shrunk, dumps(shrunk).encode("utf-8")
584
+ if rep["dropped"]:
585
+ self.log(f"code-confidential: the request was {rep['before'] // 1024} KB; {rep['dropped']} older screenshot(s) left out to send it")
580
586
  conf = self.models.get(body.get("model")) if isinstance(body, dict) and isinstance(body.get("model"), str) else None
581
587
  if not conf:
582
588
  await self._pass_through(req, res, body, target)
@@ -0,0 +1,125 @@
1
+ """OLDER IMAGES ARE SHRUNK BEFORE A REQUEST LEAVES THIS COMPUTER.
2
+
3
+ OpenCode sends the whole conversation on every call, and every screenshot in it rides along as a full base64 PNG. A model
4
+ reads an image as a fixed number of tokens whatever its bytes, so the context stays small while the REQUEST grows with each
5
+ screenshot — and TrustedRouter drops a request of about 10 MB without answering (measured 2026-10-12: 9 MB answered,
6
+ 10 MB closed the connection; a session with 18 screenshots reached 11 MB and every message in it failed).
7
+
8
+ So: the newest KEEP_FULL images go as they are; every older one is re-encoded as a JPEG no larger than MAX_SIDE on its
9
+ longest side. A screenshot of ~1 MB becomes ~100–200 KB and still reads. The result is cached by content, so the same old
10
+ screenshot shrinks to the same bytes on every call (a provider's prompt cache keeps working from there on). If the request
11
+ is STILL over BUDGET, the oldest images are replaced by a one-line note, newest kept — a request that cannot be sent helps
12
+ no one. No image library, or an image it cannot decode: that image is left as it is (the budget rule still applies).
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import base64
18
+ import hashlib
19
+ import io
20
+ import json
21
+ import re
22
+ from collections import OrderedDict
23
+ from typing import Any
24
+
25
+ KEEP_FULL = 3
26
+ MAX_SIDE = 1280
27
+ JPEG_QUALITY = 70
28
+ BUDGET = 8 * 1024 * 1024 # bytes of request body; TrustedRouter's limit is ~10 MB
29
+ CACHE_MAX = 256
30
+ _DATA_URL = re.compile(r"^data:([^;,]+)?(?:;[^,]*)?;base64,(.*)$", re.S)
31
+
32
+ _cache: "OrderedDict[str, str | None]" = OrderedDict()
33
+
34
+
35
+ def _image_parts(body: dict) -> list[tuple[int, int]]:
36
+ """(message, part) of every inline image, oldest first."""
37
+ out = []
38
+ for mi, msg in enumerate(body.get("messages") or []):
39
+ content = msg.get("content") if isinstance(msg, dict) else None
40
+ if not isinstance(content, list):
41
+ continue
42
+ for pi, part in enumerate(content):
43
+ if isinstance(part, dict) and part.get("type") == "image_url" and _DATA_URL.match(str((part.get("image_url") or {}).get("url") or "")):
44
+ out.append((mi, pi))
45
+ return out
46
+
47
+
48
+ def shrink_data_url(url: str) -> str | None:
49
+ """A smaller JPEG data URL for this image, or None (no library, cannot decode, or not smaller)."""
50
+ key = hashlib.sha256(url.encode("utf-8", "replace")).hexdigest()
51
+ if key in _cache:
52
+ _cache.move_to_end(key)
53
+ return _cache[key]
54
+ out = None
55
+ try:
56
+ from PIL import Image # optional: without it, images go as they are
57
+ m = _DATA_URL.match(url)
58
+ raw = base64.b64decode(m.group(2), validate=False)
59
+ with Image.open(io.BytesIO(raw)) as im:
60
+ im.load()
61
+ im = im.convert("RGBA") if im.mode in ("P", "LA") else im
62
+ if im.mode == "RGBA":
63
+ bg = Image.new("RGB", im.size, (255, 255, 255))
64
+ bg.paste(im, mask=im.split()[3])
65
+ im = bg
66
+ elif im.mode != "RGB":
67
+ im = im.convert("RGB")
68
+ if max(im.size) > MAX_SIDE:
69
+ im.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)
70
+ buf = io.BytesIO()
71
+ im.save(buf, "JPEG", quality=JPEG_QUALITY, optimize=True)
72
+ small = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("ascii")
73
+ out = small if len(small) < len(url) else None
74
+ except Exception: # noqa: BLE001 — an image we cannot shrink is sent as it is
75
+ out = None
76
+ _cache[key] = out
77
+ while len(_cache) > CACHE_MAX:
78
+ _cache.popitem(last=False)
79
+ return out
80
+
81
+
82
+ def _size(body: dict) -> int:
83
+ return len(json.dumps(body, separators=(",", ":"), ensure_ascii=False).encode("utf-8"))
84
+
85
+
86
+ def shrink_older_images(body: Any, keep_full: int = KEEP_FULL, budget: int = BUDGET) -> tuple[Any, dict]:
87
+ """(body, report). The body is a NEW dict when anything changed; the input is never mutated."""
88
+ report = {"shrunk": 0, "dropped": 0, "before": 0, "after": 0}
89
+ if not isinstance(body, dict):
90
+ return body, report
91
+ parts = _image_parts(body)
92
+ if not parts:
93
+ return body, report
94
+ messages = list(body["messages"])
95
+ touched: set[int] = set()
96
+
97
+ def put(mi: int, pi: int, part: dict) -> None:
98
+ if mi not in touched:
99
+ messages[mi] = {**messages[mi], "content": list(messages[mi]["content"])}
100
+ touched.add(mi)
101
+ messages[mi]["content"][pi] = part
102
+
103
+ older = parts[:-keep_full] if keep_full > 0 else parts
104
+ for mi, pi in older:
105
+ part = messages[mi]["content"][pi]
106
+ small = shrink_data_url(part["image_url"]["url"])
107
+ if small:
108
+ put(mi, pi, {**part, "image_url": {**part["image_url"], "url": small}})
109
+ report["shrunk"] += 1
110
+ out = {**body, "messages": messages} if touched else body
111
+ size = _size(out)
112
+ if size > budget:
113
+ report["before"] = size
114
+ for mi, pi in parts[:-1]: # the newest image is never dropped
115
+ if size <= budget:
116
+ break
117
+ part = messages[mi]["content"][pi]
118
+ was = len(json.dumps(part, separators=(",", ":"), ensure_ascii=False).encode("utf-8"))
119
+ note = {"type": "text", "text": "[An earlier screenshot was removed here to keep the request small enough to send.]"}
120
+ put(mi, pi, note)
121
+ size -= was - len(json.dumps(note, separators=(",", ":")).encode("utf-8"))
122
+ report["dropped"] += 1
123
+ out = {**body, "messages": messages}
124
+ report["after"] = size
125
+ return out, report
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: witbitz-code
3
- Version: 1.6.9
3
+ Version: 1.6.10
4
4
  Summary: Reach OpenCode on this computer from the Witbitz Spaces Code section, through an end-to-end sealed relay.
5
5
  Author: Witbitz
6
6
  License-Expression: MIT
@@ -19,6 +19,7 @@ Requires-Dist: websockets>=13
19
19
  Requires-Dist: httpx>=0.25
20
20
  Requires-Dist: segno>=1.5
21
21
  Requires-Dist: tinfoil<0.15,>=0.14
22
+ Requires-Dist: pillow>=10
22
23
  Provides-Extra: test
23
24
  Requires-Dist: pytest>=8; extra == "test"
24
25
  Requires-Dist: pyflakes>=3; extra == "test"
@@ -21,6 +21,7 @@ src/witbitz_code/devicelink.py
21
21
  src/witbitz_code/dirs.py
22
22
  src/witbitz_code/folders.py
23
23
  src/witbitz_code/gateway.py
24
+ src/witbitz_code/images.py
24
25
  src/witbitz_code/keys.py
25
26
  src/witbitz_code/link.py
26
27
  src/witbitz_code/mkdir.py
@@ -77,6 +78,7 @@ tests/test_devicelink.py
77
78
  tests/test_dirs.py
78
79
  tests/test_files_and_compress.py
79
80
  tests/test_folders.py
81
+ tests/test_images.py
80
82
  tests/test_outputs.py
81
83
  tests/test_pair_flows.py
82
84
  tests/test_plugins.py
@@ -3,6 +3,7 @@ websockets>=13
3
3
  httpx>=0.25
4
4
  segno>=1.5
5
5
  tinfoil<0.15,>=0.14
6
+ pillow>=10
6
7
 
7
8
  [test]
8
9
  pytest>=8
@@ -0,0 +1,102 @@
1
+ """Older screenshots are shrunk before a request leaves (images.py) — a session with 18 full PNGs reached 11 MB and
2
+ TrustedRouter dropped every call in it without answering."""
3
+
4
+ import base64
5
+ import io
6
+ import json
7
+ import os
8
+
9
+ from PIL import Image
10
+
11
+ from witbitz_code import images
12
+ from test_confidential import needs_fixture, with_proxy
13
+
14
+
15
+ def png_url(seed: int, side: int = 1800) -> str:
16
+ """A noisy PNG (noise does not compress, like a real screenshot full of text)."""
17
+ im = Image.frombytes("RGB", (side, side * 2 // 3), os.urandom(side * (side * 2 // 3) * 3) if seed >= 0 else b"\0" * (side * (side * 2 // 3) * 3))
18
+ buf = io.BytesIO()
19
+ im.save(buf, "PNG")
20
+ return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
21
+
22
+
23
+ def body_with(urls):
24
+ msgs = [{"role": "system", "content": "be useful"}]
25
+ for i, u in enumerate(urls):
26
+ msgs.append({"role": "user", "content": [{"type": "text", "text": f"shot {i}"}, {"type": "image_url", "image_url": {"url": u}}]})
27
+ msgs.append({"role": "user", "content": "Hi"})
28
+ return {"model": "openai/gpt-6-sol", "messages": msgs, "stream": True}
29
+
30
+
31
+ def urls_of(body):
32
+ return [p["image_url"]["url"] for m in body["messages"] if isinstance(m.get("content"), list) for p in m["content"] if p.get("type") == "image_url"]
33
+
34
+
35
+ def test_the_newest_images_go_as_they_are_and_older_ones_become_small_jpegs():
36
+ urls = [png_url(i, 900) for i in range(5)]
37
+ body = body_with(urls)
38
+ out, rep = images.shrink_older_images(body, keep_full=3)
39
+ got = urls_of(out)
40
+ assert got[2:] == urls[2:], "the newest three are untouched"
41
+ assert all(u.startswith("data:image/jpeg;base64,") for u in got[:2]) and rep["shrunk"] == 2
42
+ assert all(len(a) < len(b) for a, b in zip(got[:2], urls[:2]))
43
+ assert urls_of(body) == urls, "the input is never changed"
44
+ Image.open(io.BytesIO(base64.b64decode(got[0].split(",", 1)[1]))).verify()
45
+
46
+
47
+ def test_a_large_screenshot_is_brought_down_to_max_side():
48
+ out, _ = images.shrink_older_images(body_with([png_url(1, 2400), png_url(2, 300)]), keep_full=1)
49
+ with Image.open(io.BytesIO(base64.b64decode(urls_of(out)[0].split(",", 1)[1]))) as im:
50
+ assert max(im.size) == images.MAX_SIDE
51
+
52
+
53
+ def test_the_same_old_image_shrinks_to_the_same_bytes_every_call():
54
+ urls = [png_url(i, 700) for i in range(4)]
55
+ a, _ = images.shrink_older_images(body_with(urls), keep_full=1)
56
+ images._cache.clear()
57
+ b, _ = images.shrink_older_images(body_with(urls), keep_full=1)
58
+ assert urls_of(a) == urls_of(b), "deterministic, so a provider's prompt cache keeps working"
59
+
60
+
61
+ def test_no_images_or_only_new_ones_leave_the_body_as_it_was():
62
+ plain = {"model": "m", "messages": [{"role": "user", "content": "hi"}]}
63
+ assert images.shrink_older_images(plain)[0] is plain
64
+ few = body_with([png_url(1, 300)])
65
+ assert images.shrink_older_images(few)[0] is few
66
+
67
+
68
+ def test_an_image_it_cannot_decode_is_sent_as_it_is():
69
+ bad = "data:image/png;base64," + base64.b64encode(b"not an image" * 50).decode()
70
+ body = body_with([bad, png_url(1, 300)])
71
+ out, rep = images.shrink_older_images(body, keep_full=1)
72
+ assert urls_of(out)[0] == bad and rep["shrunk"] == 0
73
+
74
+
75
+ def test_over_the_budget_the_oldest_images_become_a_note_and_the_newest_stays():
76
+ urls = [png_url(i, 900) for i in range(4)]
77
+ body = body_with(urls)
78
+ out, rep = images.shrink_older_images(body, keep_full=4, budget=len(urls[0]) * 2)
79
+ assert rep["dropped"] >= 1 and urls_of(out)[-1] == urls[-1]
80
+ assert len(json.dumps(out)) <= len(urls[0]) * 2 + 2000
81
+ assert any(p.get("type") == "text" and "screenshot was removed" in p["text"] for m in out["messages"] if isinstance(m["content"], list) for p in m["content"])
82
+
83
+
84
+ def test_the_owners_session_shape_fits_under_trustedrouters_limit():
85
+ """18 screenshots of ~0.6 MB: 11 MB before; now well under the 10 MB TrustedRouter takes."""
86
+ urls = [png_url(i, 1100) for i in range(18)]
87
+ body = body_with(urls)
88
+ before = len(json.dumps(body))
89
+ out, _ = images.shrink_older_images(body)
90
+ after = len(json.dumps(out))
91
+ assert before > 10 * 1024 * 1024 and after < images.BUDGET, (before, after)
92
+
93
+
94
+ @needs_fixture
95
+ def test_the_proxy_sends_the_shrunk_body_upstream():
96
+ urls = [png_url(i, 900) for i in range(5)]
97
+
98
+ async def t(post, net, _p):
99
+ await post(body_with(urls))
100
+ sent = urls_of(json.loads(net.calls[0]["body"]))
101
+ assert sent[2:] == urls[2:] and all(u.startswith("data:image/jpeg") for u in sent[:2])
102
+ with_proxy(t)
File without changes
File without changes
File without changes