diff --git a/designer/main.cpp b/designer/main.cpp
index 08bd251..af29681 100644
--- a/designer/main.cpp
+++ b/designer/main.cpp
@@ -262,6 +262,9 @@ int main(int argc, char* argv[]) {
         } else if (arg == "--playback" && i + 1 < argc) {
             g_test_mode = TestMode::Playback;
             g_test_input_path = argv[++i];
+        } else if (arg == "--screenshot-after" && i + 1 < argc) {
+            // Capture screenshot after playback completes (used with --playback)
+            g_test_output_path = argv[++i];
         } else if (arg == "--screenshot" && i + 1 < argc) {
             g_test_mode = TestMode::Screenshot;
             g_test_output_path = argv[++i];
@@ -515,7 +518,20 @@ int main(int argc, char* argv[]) {
             // Check if playback finished
             if (g_action_player->IsFinished()) {
                 std::cout << "Playback complete" << std::endl;
-                // Optionally exit after playback
+
+                // Capture screenshot if --screenshot-after was specified
+                if (!g_test_output_path.empty()) {
+                    int fb_width, fb_height;
+                    glfwGetFramebufferSize(g_window, &fb_width, &fb_height);
+                    mosis::testing::VisualCapture capture(fb_width, fb_height);
+                    if (capture.CaptureScreenshot(g_test_output_path)) {
+                        std::cout << "Screenshot saved to: " << g_test_output_path << std::endl;
+                    } else {
+                        std::cerr << "Failed to save screenshot" << std::endl;
+                    }
+                }
+
+                // Exit after playback
                 glfwSetWindowShouldClose(g_window, GLFW_TRUE);
             }
         }
diff --git a/test-apps/com.mosis.sandbox-test/app.lua b/test-apps/com.mosis.sandbox-test/app.lua
index 9f81bc9..264a92a 100644
--- a/test-apps/com.mosis.sandbox-test/app.lua
+++ b/test-apps/com.mosis.sandbox-test/app.lua
@@ -3,13 +3,41 @@
 
 local results = {}
 local logCounter = 0
+local currentDocument = nil  -- Cache the document reference
+
+-- Helper to get document (try multiple methods)
+local function getDocument()
+    -- First try the cached document
+    if currentDocument then
+        return currentDocument
+    end
+    -- Try the global document
+    if document then
+        currentDocument = document
+        return document
+    end
+    -- Try to get from RmlUi context
+    if rmlui and rmlui.contexts and rmlui.contexts.main then
+        local ctx = rmlui.contexts.main
+        -- documents is a proxy, iterate to get first doc
+        if ctx.documents then
+            for i, doc in ipairs(ctx.documents) do
+                if doc then
+                    currentDocument = doc
+                    return currentDocument
+                end
+            end
+        end
+    end
+    return nil
+end
 
 local function log(msg)
     logCounter = logCounter + 1
     table.insert(results, string.format("[%03d] %s", logCounter, msg))
-    -- document may not be available during initial script load
-    if document then
-        local el = document:GetElementById("results")
+    local doc = getDocument()
+    if doc then
+        local el = doc:GetElementById("results")
         if el then
             el.inner_rml = table.concat(results, "\n")
         end
@@ -26,20 +54,26 @@ function goBack()
 end
 
 local function setStatus(id, status, success)
-    -- document may not be available during initial script load
-    if not document then return end
-    local el = document:GetElementById(id)
+    local doc = getDocument()
+    if not doc then
+        print("[LUA] ERROR: document not available!")
+        return
+    end
+    local el = doc:GetElementById(id)
     if el then
         if success then
             el.inner_rml = "&#x2713; " .. status
         else
             el.inner_rml = "&#x2717; " .. status
         end
+    else
+        print("[LUA] ERROR: element not found: " .. id)
     end
 end
 
 -- Timer test
 function testTimer()
+    print("[LUA] testTimer called!")
     setStatus("timer-status", "Running...", true)
     log("Starting timer test...")
 
@@ -206,6 +240,10 @@ function testStorage()
     end
 end
 
--- Initialize
+-- Initialize: cache the document when script loads
+if document then
+    currentDocument = document
+    print("[LUA] Document cached on load")
+end
 log("Sandbox Test App loaded")
 log("Lua version: " .. (_VERSION or "unknown"))