diff --git a/priv/python/albedo_api.py b/priv/python/albedo_api.py index d9f929ce..da5575bf 100644 --- a/priv/python/albedo_api.py +++ b/priv/python/albedo_api.py @@ -171,8 +171,8 @@ class PythonApi: # search job, say) never crowds the model's own output out of retention. forget_output: Callable[[str], None] | None = None # Attaches PNG, JPEG, or WebP bytes to the running cell's result, so the - # model sees the image; returns a short description, raises ValueError - # past the per-cell limits. Absent when the host pre-dates images. + # model sees the image; returns a description or an overflow file path. + # Invalid images raise ValueError. Absent when the host pre-dates images. attach_image: Callable[[bytes], str] | None = None # A host call that blocks until its answer, for synchronous APIs over # routes that answer at once. Absent when the host pre-dates it. diff --git a/priv/python/albedo_capture.py b/priv/python/albedo_capture.py index 639cef04..c451a771 100644 --- a/priv/python/albedo_capture.py +++ b/priv/python/albedo_capture.py @@ -201,8 +201,13 @@ class Capture: return self.read(0, PREVIEW) return self.tail().decode("utf-8", errors="ignore") - def attach(self, data: bytes) -> str: - """Queue one image for this cell's result; raises ValueError past the limits.""" + def attach(self, data: bytes | os.PathLike[str]) -> str: + """Queue one image, or save overflow that would fit in an empty cell.""" + if isinstance(data, os.PathLike): + with open(data, "rb") as source: + self._check_image_bytes(os.fstat(source.fileno()).st_size) + data = source.read(MAX_IMAGE_BYTES + 1) + self._check_image_bytes(len(data)) mime = image_type(data) if mime is None: raise ValueError("image must be PNG, JPEG, or WebP bytes") @@ -219,15 +224,34 @@ class Capture: raise ValueError( f"{width}x{height} image is over this model's {self.max_edge}px edge limit" ) - if len(self.images) >= MAX_IMAGES: - raise ValueError(f"a cell returns at most {MAX_IMAGES} images") - if sum(map(len, self.images)) + len(data) > MAX_IMAGE_BYTES: + if ( + len(self.images) >= MAX_IMAGES + or sum(map(len, self.images)) + len(data) > MAX_IMAGE_BYTES + ): + return self._save_image(data, mime) + self.images.append(data) + return f"attached {mime}, {len(data)} bytes" + + def _check_image_bytes(self, size: int) -> None: + if size > MAX_IMAGE_BYTES: raise ValueError( - f"a cell's images total at most {MAX_IMAGE_BYTES} bytes; " - f"this {len(data)}-byte image does not fit" + f"this {size}-byte image exceeds the 5 MiB ({MAX_IMAGE_BYTES} bytes) " + "per-cell limit; downscale or crop it before showing it" ) - self.images.append(data) - return f"{mime}, {len(data)} bytes" + + def _save_image(self, data: bytes, mime: str) -> str: + """Save an overflow image without adding it to this cell's result.""" + suffix = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}[ + mime + ] + fd, path = tempfile.mkstemp(prefix="image-", suffix=suffix, dir=SPILL_DIR) + try: + with os.fdopen(fd, "wb") as destination: + destination.write(data) + except BaseException: + os.unlink(path) + raise + return f"not attached, this cell's image count or 5 MiB budget is full: saved to {path}; show_image({path!r}) in another cell" def encoded_images(self) -> list[str]: return [base64.b64encode(image).decode("ascii") for image in self.images] diff --git a/priv/python/albedo_cells.py b/priv/python/albedo_cells.py index 93d3a6fe..2660ccae 100644 --- a/priv/python/albedo_cells.py +++ b/priv/python/albedo_cells.py @@ -7,7 +7,7 @@ import albedo_bundle import albedo_state from albedo_values import Unencodable, decode, encode import albedo_capture -from albedo_capture import Capture, PREVIEW, MAX_IMAGE_BYTES +from albedo_capture import Capture, PREVIEW import albedo_link import albedo_memory from albedo_protocol import parse_incoming @@ -363,32 +363,41 @@ def host_now(method: str, args: dict[str, object]) -> object: def attach_image(data: bytes) -> str: """Attach image bytes to the running cell; plugins reach this as api.attach_image.""" + return _image_capture().attach(bytes(data)) + + +def _image_capture() -> Capture: capture = CELL.get() # Spawned tasks inherit context even after their creating cell finishes. if capture is None or capture.kind != "cell" or capture.id not in ACTIVE_CAPTURES: raise RuntimeError( "images attach to a running cell; background work has no result to carry them" ) - return capture.attach(bytes(data)) + return capture def show_image( source: bytes | bytearray | memoryview | str | os.PathLike[str], + *sources: bytes | bytearray | memoryview | str | os.PathLike[str], ) -> albedo_api.Text: - """Return an image to yourself with this cell's result. `source` is PNG, - JPEG, or WebP bytes, or a path to such a file. At most 4 images and 5 MiB - per cell, each within the model's edge limit, or it raises ValueError; the - image arrives after the cell finishes, not mid-cell.""" - if isinstance(source, (bytes, bytearray, memoryview)): - data = bytes(source) - elif isinstance(source, (str, os.PathLike)): - with open(source, "rb") as file: - data = file.read(MAX_IMAGE_BYTES + 1) - else: - raise TypeError( - f"show_image takes image bytes or a path, not {type(source).__name__}" - ) - return albedo_api.Text("attached " + attach_image(data)) + """Show PNG/JPEG/WebP bytes or paths, several at once, after the cell finishes. + + At most 4 images and 5 MiB attach per cell, within the model's edge limit. + Overflow is saved to paths named in the result; a single image over 5 MiB raises. + """ + capture = _image_capture() + results = [] + for item in (source, *sources): + if isinstance(item, (bytes, bytearray, memoryview)): + data = bytes(item) + elif isinstance(item, (str, os.PathLike)): + data = pathlib.Path(item) + else: + raise TypeError( + f"show_image takes image bytes or a path, not {type(item).__name__}" + ) + results.append(capture.attach(data)) + return albedo_api.Text("\n".join(results)) def display( diff --git a/priv/python/albedo_plugins/remote.py b/priv/python/albedo_plugins/remote.py index 61259008..f09ca289 100644 --- a/priv/python/albedo_plugins/remote.py +++ b/priv/python/albedo_plugins/remote.py @@ -1026,7 +1026,7 @@ class RemoteConnection: raise TypeError( f"show_image takes image bytes or a remote path, not {type(source).__name__}" ) - return "attached " + attach_image(data) + return attach_image(data) async def _ssh( self, script: str, timeout: float, stdin: bytes | None = None diff --git a/robot-docs/kernel.md b/robot-docs/kernel.md index a8f96b2e..eea16eb6 100644 --- a/robot-docs/kernel.md +++ b/robot-docs/kernel.md @@ -37,6 +37,19 @@ later cells. The process-wide `builtins` module is unchanged. Plugin instances are owned by this kernel's API, not module-level singletons. This fallback is built once after plugins load and is not serialized with user state. +## images + +`show_image(source, *sources)` returns PNG, JPEG, or WebP bytes or paths with the +cell result. `albedo_capture.Capture` owns the four-image and 5 MiB attachment +budgets. Overflow is saved losslessly under the session's output directory (or +in a temporary directory without one); the returned Text names each file and +says to use `show_image(path)` in another cell. Output-directory files share +its 14-day retention. An image over 5 MiB cannot fit even in an empty cell: it +raises with its size and advice to downscale or crop, without reading a large +file path into memory. +Invalid, cut-short, or over-edge attachments still raise. No image decoder is +required, so overflow is not tiled. + ## spawned browsers `priv/python/albedo_plugins/browser/launch.py` owns browser discovery and spawned diff --git a/src/albedo/harness/extensions/python/extension.gleam b/src/albedo/harness/extensions/python/extension.gleam index c8918205..1eab8ce1 100644 --- a/src/albedo/harness/extensions/python/extension.gleam +++ b/src/albedo/harness/extensions/python/extension.gleam @@ -182,7 +182,7 @@ pub fn extension() -> extension.Extension { ]), extension.ContextPlugin(place.context), extension.ToolPlugin( - "Every capability this session has is exposed as a Python binding inside the python kernel: there are no separate tool calls, so call these names inside a python cell. Python cells automatically background after 60 seconds; the session wakes when they finish, so do not poll or sleep. Independent work starts together, not one after another: await parallel(a, b, ...) awaits several calls or run() handles at once and returns their results in order (return_exceptions=True keeps the successes); independent children spawn back to back. Await only when the next step needs the result. Python has a persistent namespace, top-level await, cells.read/info/trace and cells.run for saved-source repair (all async), and output.read/output.list for bounded retained output (synchronous; awaiting them also works). cells.last_id is the id of the latest cell, and await cells.list(limit=20) lists this session's cells newest first with their status and first line, even after their output has rolled out. output.read(id, offset=0, limit=4000) returns up to limit characters, at most 65536 per call (page with offset); offset and limit may be positional, and job.tail()/job.head() answer awaitably too. output.list() names every retained channel, which is also how to find earlier cells: cells, background jobs, and 'native' for bytes written to fd 1/2 while no cell was running. show_image(source) returns a PNG, JPEG, or WebP (bytes or a file path) to you with this cell's result, so you see it after the cell ends; at most 4 images and 5 MiB per cell. display(content, title='', description='') shows text, code, or diffs directly to the user in the transcript (folding long blocks); use display() to display any diffs, text, or files you want the user to see instead of regurgitating them in your output to save tokens.", + "Every capability this session has is exposed as a Python binding inside the python kernel: there are no separate tool calls, so call these names inside a python cell. Python cells automatically background after 60 seconds; the session wakes when they finish, so do not poll or sleep. Independent work starts together, not one after another: await parallel(a, b, ...) awaits several calls or run() handles at once and returns their results in order (return_exceptions=True keeps the successes); independent children spawn back to back. Await only when the next step needs the result. Python has a persistent namespace, top-level await, cells.read/info/trace and cells.run for saved-source repair (all async), and output.read/output.list for bounded retained output (synchronous; awaiting them also works). cells.last_id is the id of the latest cell, and await cells.list(limit=20) lists this session's cells newest first with their status and first line, even after their output has rolled out. output.read(id, offset=0, limit=4000) returns up to limit characters, at most 65536 per call (page with offset); offset and limit may be positional, and job.tail()/job.head() answer awaitably too. output.list() names every retained channel, which is also how to find earlier cells: cells, background jobs, and 'native' for bytes written to fd 1/2 while no cell was running. show_image(source, *sources) is the default way to return PNG, JPEG, or WebP images (bytes or file paths) with this cell's result; at most 4 images and 5 MiB attach per cell. Overflow images that fit an empty cell are saved, and the returned Text names each path: use show_image(path) in another cell. A single image over 5 MiB raises; downscale or crop it. display(content, title='', description='') shows text, code, or diffs directly to the user in the transcript (folding long blocks); use display() to display any diffs, text, or files you want the user to see instead of regurgitating them in your output to save tokens.", [extension.Tool(definition(), invoke, recover)], [], [ diff --git a/test/e2e/image_limits_test.py b/test/e2e/image_limits_test.py index 45d4c3dc..2864f04c 100644 --- a/test/e2e/image_limits_test.py +++ b/test/e2e/image_limits_test.py @@ -61,6 +61,103 @@ def cell_result(output): class ImageLimitsTest(unittest.TestCase): + def test_overflow_images_are_saved_and_can_be_shown_in_another_cell(self): + small = base64.b64encode(png(2, 3)).decode() + code = f""" +image = base64.b64decode({small!r}) +notice = str(show_image(image, image, image, image, image, image)) +print(notice) +paths = re.findall(r"saved to (.*?); show_image", notice) +assert len(paths) == 2, notice +assert all(Path(path).read_bytes() == image for path in paths) +print('six images preserved') +""" + calls = 0 + + def reply(request): + nonlocal calls + calls += 1 + if calls == 1: + return python(code) + if calls == 2: + return python("show_image(*paths)") + return text("done") + + provider = Provider(reply) + self.addCleanup(provider.close) + with Albedo(provider, protocol="responses") as app: + session = app.session() + app.prompt(session, "show six images").close() + app.idle(session) + outputs = [ + item + for item in provider.requests[-1]["request"]["input"] + if item.get("type") == "function_call_output" + ] + self.assertEqual(len(outputs), 2) + results = [cell_result(item) for item in outputs] + self.assertTrue(all(r["status"] == "ok" for r in results), results) + self.assertIn("six images preserved", results[0]["output"]) + self.assertIn("in another cell", results[0]["output"]) + counts = [ + len(re.findall("data:image/png;base64,", json.dumps(item))) + for item in outputs + ] + self.assertEqual(counts, [4, 2]) + + def test_image_byte_budget_saves_overflow_without_failing_the_cell(self): + import inspect + + code = ( + "import struct, zlib\n" + + inspect.getsource(png) + + """ +image = png(2, 3, b'x' * (3 * 1024 * 1024)) +notice = str(show_image(image, image)) +print(notice) +path = re.search(r"saved to (.*?); show_image", notice)[1] +assert Path(path).read_bytes() == image +large = png(2, 3, b'x' * (6 * 1024 * 1024)) +original = Path('large.png') +original.write_bytes(large) +for source in (original, large): + try: + show_image(source) + except ValueError as error: + message = str(error) + assert str(len(large)) in message, message + assert '5 MiB' in message, message + assert 'downscale or crop' in message, message + else: + raise AssertionError('oversize image must not suggest showing it again') +assert original.read_bytes() == large +print('byte overflow preserved') +""" + ) + + def reply(request): + if request["input"][-1].get("type") == "function_call_output": + return text("done") + return python(code) + + provider = Provider(reply) + self.addCleanup(provider.close) + with Albedo(provider, protocol="responses") as app: + session = app.session() + app.prompt(session, "show large images").close() + app.idle(session) + output = next( + item + for item in provider.requests[-1]["request"]["input"] + if item.get("type") == "function_call_output" + ) + result = cell_result(output) + self.assertEqual(result["status"], "ok", result) + self.assertIn("byte overflow preserved", result["output"]) + self.assertEqual( + len(re.findall("data:image/png;base64,", json.dumps(output))), 1 + ) + @exclusive def test_large_valid_image_remains_available_after_durable_reload(self): # A valid ancillary chunk makes encoded length exceed 5 MiB while the