Feat/screenshot system improvements #864
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Remove preserveDrawingBuffer requirement by forcing render before capture - Add ScreenshotProvider component that registers capture function in store - Download screenshots directly to client browser (no server upload for UI) - Add FrontendSession.take_screenshot() for programmatic screenshot capture - Add POST /api/sessions/<id>/screenshot endpoint for Python clients - Add take_screenshot MCP tool for AI-assisted screenshot capture - Support path-tracing mode (skip render to preserve accumulated samples) - Remove preserve_drawing_buffer field from Camera model - Clean up dead code (captureCanvasScreenshot, uploadScreenshot, etc.) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Root cause: The @react-three/gpu-pathtracer library has a bug where it resets on every re-render due to rest props spread creating new object references in its useLayoutEffect dependencies. Fixes applied: - Wrap MyScene in React.memo to prevent re-renders from parent - Fix useFrameLoadTime hook to use individual selectors instead of subscribing to entire store (was causing re-renders on any state change) - Move snackbar rendering from Canvas.tsx to landingPage.tsx to avoid snackbar state changes triggering Canvas re-renders - Move screenshot button from dropdown menu to header bar - Add PathtracingCaptureProvider for pathtracing screenshot support Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR refactors screenshot capture and upload functionality across the frontend and backend. On the frontend, it moves screenshot capture from RoomManagementMenu to the landing page, introduces dedicated providers (ScreenshotProvider and PathtracingCaptureProvider) for centralized screenshot management with pathtracer support, and implements client-side download capability. On the backend, it adds REST API endpoints for screenshot requests with Redis-backed status tracking and Socket.IO signaling. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User (Landing Page)
participant Frontend as Frontend App
participant ScreenshotProvider
participant Pathtracer as Pathtracer/<br/>Canvas
participant ScreenshotCapture as PathtracingCapture<br/>Provider
participant Backend as Backend API
User->>Frontend: Click "Take Screenshot" Button
Frontend->>ScreenshotProvider: Call getCapture()
alt Pathtracer Available
ScreenshotProvider->>ScreenshotCapture: Use pathtracerCapture
ScreenshotCapture->>Pathtracer: Capture Canvas (gl.domElement)
Pathtracer-->>ScreenshotCapture: PNG Blob (with preserveDrawingBuffer)
ScreenshotCapture-->>ScreenshotProvider: Return Blob Promise
else Standard Fallback
ScreenshotProvider->>Canvas: Force render + Capture
Canvas-->>ScreenshotProvider: Return PNG Blob
end
ScreenshotProvider-->>Frontend: Blob received
Frontend->>Frontend: downloadScreenshot(blob)
Frontend-->>User: Browser downloads PNG file
Frontend->>Frontend: Show Success Snackbar
sequenceDiagram
participant Client as Frontend Client
participant Backend as Backend API
participant Redis as Redis Store
participant SocketIO as Socket.IO Event
participant Session as Session/<br/>Room
Client->>Backend: POST /api/sessions/{id}/screenshot<br/>(timeout=10s)
Backend->>Backend: Validate user membership<br/>Generate request_id
Backend->>Redis: Store pending state<br/>(TTL: 120s)
Backend->>SocketIO: Emit screenshot:request<br/>(request_id, uploadUrl)
SocketIO->>Session: Notify connected session
Session->>Session: Capture screenshot via<br/>ScreenshotProvider
Session->>Backend: POST upload with request_id
Backend->>Redis: Write completion status<br/>(TTL: 60s)
Backend-->>Redis: Poll for completion<br/>(100ms intervals)
Redis-->>Backend: Status: completed
Backend-->>Client: Return 200<br/>(screenshot_id, room_id)
Client->>Client: Show Success Message
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
app/src/components/PathTracingRenderer.tsx (1)
1-10: Fix TypeScript compilation error:Reactnamespace not in scope.The code uses
React.ReactNodewithout importing React. With"react-jsx"in the tsconfig, JSX transforms automatically, but type annotations still require explicit React imports. This will cause TS compilation to fail.Add either a React import (like in FigureWindow.tsx) or import ReactNode directly:
Proposed fix
+import type { ReactNode } from "react"; import { Pathtracer } from "@react-three/gpu-pathtracer"; import { Environment } from "@react-three/drei"; import type { PathTracing } from "../types/room-config"; import { PathtracingUpdater } from "./PathtracingUpdater"; import { PathtracingCaptureProvider } from "./three/PathtracingScreenshotCapture"; interface PathTracingRendererProps { settings?: PathTracing; - children: React.ReactNode; + children: ReactNode; }
🤖 Fix all issues with AI agents
In `@src/zndraw-mcp/__init__.py`:
- Around line 87-175: The function take_screenshot erroneously requires a local
server via find_running_server() even when an explicit url is provided; change
the logic so you only call/find_running_server() (and error if not found) when
url is None, otherwise skip the local-server check and proceed to construct
ZnDraw(url=url, ... ) and continue; update the conditional around
server_info/find_running_server to respect the url parameter and maintain the
rest of the session selection and screenshot handling in take_screenshot
unchanged.
In `@src/zndraw/app/session_routes.py`:
- Around line 223-325: The request_screenshot route allows any authenticated
user to trigger screenshots for any session_id; after you retrieve session_data
and room_id in request_screenshot, import and call get_current_user (from
zndraw.auth) and verify the current user is a member of that room_id (e.g.,
check room_id in the user's membership/rooms collection); if the user is not a
member, return a 403 error before generating request_id or setting the Redis
request key. Ensure this check is placed just after room_id is obtained and
before the Redis setex/socketio emit logic so only room members can request
screenshots.
In `@src/zndraw/session_manager.py`:
- Around line 186-221: The docstring for take_screenshot omits that the
underlying API calls can raise requests.HTTPError; update the Raises section of
the take_screenshot docstring to list requests.HTTPError (thrown by
_vis.api.request_screenshot and _vis.api.download_screenshot) in addition to
TimeoutError, and add a short sentence in the description noting that HTTP
errors from those API calls are propagated to the caller.
🧹 Nitpick comments (3)
app/src/components/PathtracingUpdater.tsx (1)
31-61: Gate debug logging to avoid noisy production consoles.Consider wrapping these
console.logcalls behind a debug flag or a logger utility.app/src/components/three/ScreenshotProvider.tsx (2)
48-50: Quality parameter is ignored for PNG format.The third parameter (1.0) to
toBlob()is the quality setting, which only applies to lossy formats likeimage/jpegorimage/webp. For PNG, it's silently ignored. Consider removing it to avoid confusion, or document this behavior.🔧 Suggested fix
gl.domElement.toBlob( (blob) => { if (blob) { resolve(blob); } else { reject(new Error("Failed to create blob from canvas")); } }, "image/png", - 1.0, );
103-116: Consider adding a timeout to the upload request.The fetch call has no timeout, which could cause the request to hang indefinitely if the server is unresponsive. For better reliability, consider using
AbortControllerwith a timeout.🔧 Suggested implementation with timeout
+ const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 30000); // 30s timeout + const response = await fetch(uploadUrl, { method: "POST", body: formData, + signal: controller.signal, }); + clearTimeout(timeoutId);
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #864 +/- ##
==========================================
- Coverage 80.03% 79.78% -0.26%
==========================================
Files 165 165
Lines 20133 20213 +80
==========================================
+ Hits 16114 16127 +13
- Misses 4019 4086 +67 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- Use send_file instead of send_from_directory for screenshot download - Resolve media_path to absolute path to avoid working directory issues - Add ScreenshotImage dataclass with Jupyter display support (_repr_png_) - take_screenshot() now returns ScreenshotImage with .save() and .data Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- PathTracingRenderer.tsx: import ReactNode properly from react - session_routes.py: add room membership check in request_screenshot - session_manager.py: document requests.HTTPError in take_screenshot docstring - zndraw-mcp/__init__.py: only check for local server when url is None Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@src/zndraw-mcp/__init__.py`:
- Around line 158-167: session.take_screenshot() now returns a ScreenshotImage
object, so the current code that writes image_bytes will raise a TypeError;
update the try block that calls session.take_screenshot() to extract the raw
bytes (e.g., use the ScreenshotImage.data attribute) or call the object's save
method before writing to disk (replace writing image_bytes with writing
screenshot.data or saving via ScreenshotImage.save), keeping the same tempfile
NamedTemporaryFile flow and preserving filepath variable for downstream use.
In `@src/zndraw/app/session_routes.py`:
- Around line 278-284: Replace the ad-hoc timeout parsing with a Pydantic
request model: define a small model (e.g., class GenerateRequest(BaseModel))
with a timeout field having a default of 10 and validation to ensure it's a
float between 0 and 60 (use conint/confloat or validator to cap at 60). In the
route, call GenerateRequest.parse_obj(request.get_json(silent=True) or {})
inside a try/except catching pydantic.ValidationError, returning a 400 with the
validation message on error, and then use request_model.timeout instead of the
manual timeout variable; remove the existing request.get_json fallback and the
try/except ValueError/TypeError block. Ensure you import BaseModel and
ValidationError from pydantic.
In `@src/zndraw/session_manager.py`:
- Around line 221-259: The take_screenshot method's return type annotation is
incorrect: it returns a ScreenshotImage instance but is declared as -> bytes.
Update the signature of take_screenshot in session_manager.py to return the
correct type (ScreenshotImage) and ensure ScreenshotImage is imported or in
scope so type checkers recognize it; leave the implementation unchanged (it
already constructs and returns ScreenshotImage(data=data)).
🧹 Nitpick comments (1)
src/zndraw/app/screenshot_chat_routes.py (1)
216-262: Avoid completing non‑existent requests and preserve TTL.
request_idis client‑supplied;setexunconditionally creates/overwrites the key even if no pending request exists and resets TTL from 120s (set at request creation) to 60s. Useset()withxx=True, keepttl=Trueto update only when the request key exists and preserve its original TTL, preventing spoofed or resurrected completions.♻️ Proposed change
- # Update request status to completed with screenshot ID - request_data = { - "status": "completed", - "screenshot_id": screenshot.id, - } - # Keep the same TTL (60s should be enough for the caller to read it) - r.setex(request_key, 60, json.dumps(request_data)) + # Update request status to completed with screenshot ID + request_data = { + "status": "completed", + "screenshot_id": screenshot.id, + } + # Update only if key exists; preserve TTL + updated = r.set( + request_key, + json.dumps(request_data), + xx=True, + keepttl=True, + ) + if not updated: + log.warning( + "upload_screenshot: request %s not found or expired; skipping", + request_id, + )redis-py 6.4.0+ (as required in pyproject.toml) supports both
xxandkeepttlparameters.
src/zndraw/session_manager.py
Outdated
Comment on lines
221
to
259
| def take_screenshot(self, timeout: float = 10.0) -> bytes: | ||
| """Request a screenshot from this browser session. | ||
| Uses REST-based communication: | ||
| 1. Request triggers lightweight Socket.IO signal to browser | ||
| 2. Browser captures screenshot and uploads via REST | ||
| 3. This method downloads the result via REST | ||
| Parameters | ||
| ---------- | ||
| timeout : float | ||
| Maximum wait time in seconds (default: 10.0). | ||
| Returns | ||
| ------- | ||
| ScreenshotImage | ||
| Screenshot with Jupyter display support. | ||
| Raises | ||
| ------ | ||
| TimeoutError | ||
| If no response within timeout. | ||
| requests.HTTPError | ||
| If the server returns an error response. | ||
| Examples | ||
| -------- | ||
| >>> session = vis.sessions["abc-123"] | ||
| >>> img = session.take_screenshot() | ||
| >>> img # Displays in Jupyter | ||
| >>> img.save("screenshot.png") | ||
| """ | ||
| # Request screenshot from browser session | ||
| result = self._vis.api.request_screenshot(self.session_id, timeout) | ||
|
|
||
| # Download the actual image data | ||
| screenshot_id = result["screenshot_id"] | ||
| data = self._vis.api.download_screenshot(screenshot_id) | ||
| return ScreenshotImage(data=data) |
There was a problem hiding this comment.
Fix return type annotation to match ScreenshotImage.
Line 221: The method returns a ScreenshotImage instance, but the signature says -> bytes. This will mislead callers and type checkers.
✅ Suggested fix
- def take_screenshot(self, timeout: float = 10.0) -> bytes:
+ def take_screenshot(self, timeout: float = 10.0) -> ScreenshotImage:🤖 Prompt for AI Agents
In `@src/zndraw/session_manager.py` around lines 221 - 259, The take_screenshot
method's return type annotation is incorrect: it returns a ScreenshotImage
instance but is declared as -> bytes. Update the signature of take_screenshot in
session_manager.py to return the correct type (ScreenshotImage) and ensure
ScreenshotImage is imported or in scope so type checkers recognize it; leave the
implementation unchanged (it already constructs and returns
ScreenshotImage(data=data)).
- MCP take_screenshot: extract .data from ScreenshotImage return type - session_manager: fix return type annotation -> ScreenshotImage - session_routes: replace ad-hoc timeout parsing with Pydantic model Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/zndraw/app/session_routes.py`:
- Around line 319-335: The polling loop currently computes max_polls =
int(timeout / poll_interval) which floors the count and can prematurely stop
polling; change this to use math.ceil(timeout / poll_interval) (e.g., max_polls
= math.ceil(timeout / poll_interval)) so the loop covers the full timeout
duration, and add/ensure an import math at the top of the module; update the
code around poll_interval, max_polls, and the for _ in range(max_polls) loop in
session_routes.py accordingly.
🧹 Nitpick comments (1)
src/zndraw-mcp/__init__.py (1)
121-121: Movetempfileimport to module scope.Line 121: this local import doesn't appear to be for lazy loading; keeping imports at the top improves consistency.
♻️ Suggested change
@@ -import subprocess -import traceback +import subprocess +import tempfile +import traceback @@ - import tempfileAs per coding guidelines, keep imports at the top.
src/zndraw/app/session_routes.py
Outdated
Comment on lines
319
to
335
| # Poll for completion | ||
| poll_interval = 0.1 # 100ms | ||
| max_polls = int(timeout / poll_interval) | ||
|
|
||
| for _ in range(max_polls): | ||
| data_str = r.get(request_key) | ||
| if data_str: | ||
| data = json.loads(data_str) | ||
| if data.get("status") == "completed": | ||
| screenshot_id = data.get("screenshot_id") | ||
| # Clean up the request key | ||
| r.delete(request_key) | ||
| return { | ||
| "screenshot_id": screenshot_id, | ||
| "room_id": room_id, | ||
| }, 200 | ||
| time.sleep(poll_interval) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd src/zndraw/app && head -340 session_routes.py | tail -30Repository: zincware/ZnDraw
Length of output: 1040
🏁 Script executed:
cd src/zndraw/app && head -10 session_routes.pyRepository: zincware/ZnDraw
Length of output: 377
🏁 Script executed:
rg "math" src/zndraw/app/session_routes.pyRepository: zincware/ZnDraw
Length of output: 41
🏁 Script executed:
cd src/zndraw/app && head -30 session_routes.pyRepository: zincware/ZnDraw
Length of output: 858
🏁 Script executed:
cd src/zndraw/app && rg "^import|^from" session_routes.py | head -20Repository: zincware/ZnDraw
Length of output: 425
🏁 Script executed:
cd src/zndraw/app && sed -n '280,340p' session_routes.pyRepository: zincware/ZnDraw
Length of output: 2012
🏁 Script executed:
rg "request_screenshot" src/zndraw/app/session_routes.py -B 5 -A 50 | head -80Repository: zincware/ZnDraw
Length of output: 2567
🏁 Script executed:
cd src/zndraw/app && rg "class ScreenshotRequest" session_routes.py -A 15Repository: zincware/ZnDraw
Length of output: 449
🏁 Script executed:
cd src/zndraw/app && python3 << 'EOF'
# Test the behavior of int() vs math.ceil() for various timeout values
import math
test_cases = [
0.05, # Less than poll_interval
0.1, # Exactly one interval
0.15, # Between 1 and 2 intervals
0.25, # 2.5 intervals
0, # Zero timeout
-1, # Negative timeout
10, # Normal case
]
poll_interval = 0.1
print("timeout\t| int()\t| ceil()\t| iterations (int)\t| iterations (ceil)")
print("-" * 70)
for timeout in test_cases:
int_result = int(timeout / poll_interval)
ceil_result = math.ceil(timeout / poll_interval)
iterations_int = list(range(int_result))
iterations_ceil = list(range(ceil_result))
print(f"{timeout}\t| {int_result}\t| {ceil_result}\t| {len(iterations_int)}\t\t\t| {len(iterations_ceil)}")
EOFRepository: zincware/ZnDraw
Length of output: 145
🏁 Script executed:
cd src/zndraw/app && rg "timeout" session_routes.py -B 2 -A 2Repository: zincware/ZnDraw
Length of output: 782
Use math.ceil() instead of int() to avoid premature timeout in polling loop.
Line 321 floors the poll count using int(), which causes early timeout for non-multiple intervals (e.g., timeout=0.15s with 0.1s interval yields 1 poll instead of 2). Use math.ceil() to round up and ensure the full timeout duration.
Suggested fix
import json
import logging
+import math
import time- max_polls = int(timeout / poll_interval)
+ max_polls = math.ceil(timeout / poll_interval)🤖 Prompt for AI Agents
In `@src/zndraw/app/session_routes.py` around lines 319 - 335, The polling loop
currently computes max_polls = int(timeout / poll_interval) which floors the
count and can prematurely stop polling; change this to use math.ceil(timeout /
poll_interval) (e.g., max_polls = math.ceil(timeout / poll_interval)) so the
loop covers the full timeout duration, and add/ensure an import math at the top
of the module; update the code around poll_interval, max_polls, and the for _ in
range(max_polls) loop in session_routes.py accordingly.
Simplified documentation following KISS principle: - Removed obvious/redundant comments - Shortened verbose docstrings - Removed "DRY" and implementation detail explanations - Kept only essential context where not self-evident from code Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/zndraw/app/screenshot_chat_routes.py (1)
206-257: Move upload form defaults into a request model.
formatcurrently defaults viarequest.form.get("format", "png"); per guidelines, defaults should be centralized in a Pydantic model rather than scattered in route logic. As per coding guidelines, defaults should live in Pydantic models.♻️ Suggested refactor (model + parsing)
-from flask import Blueprint, current_app, request, send_file +from flask import Blueprint, current_app, request, send_file +from pydantic import BaseModel, Field, ValidationError @@ +class ScreenshotUpload(BaseModel): + """Validated form fields for screenshot upload.""" + format: str = Field(default="png") + width: int | None = None + height: int | None = None + request_id: str | None = None @@ - format = request.form.get("format", "png") - width = request.form.get("width", type=int) - height = request.form.get("height", type=int) - request_id = request.form.get("request_id") + try: + form = ScreenshotUpload.model_validate(request.form) + except ValidationError as e: + return {"error": e.errors()}, 400 + format = form.format + width = form.width + height = form.height + request_id = form.request_id
♻️ Duplicate comments (1)
src/zndraw/app/session_routes.py (1)
310-327: Avoid premature timeouts by rounding up poll count.Line 311 uses
int(timeout / poll_interval), which truncates and can shorten the effective wait (e.g., 0.15s → 1 poll). Usemath.ceilto cover the full timeout.🐛 Proposed fix
-import time +import math +import time @@ - max_polls = int(timeout / poll_interval) + max_polls = math.ceil(timeout / poll_interval)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.