From a7dac11aad217efb38660a85bdc2e3eaebaff0a8 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Wed, 29 Oct 2025 11:09:43 +0800
Subject: [PATCH] viz: keep rewrite step in back button history (#12986)

---
 tinygrad/viz/js/index.js | 3 ++-
 tinygrad/viz/serve.py    | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js
index 3f72816cc9..9fc3c76728 100644
--- a/tinygrad/viz/js/index.js
+++ b/tinygrad/viz/js/index.js
@@ -585,12 +585,13 @@ const evtSources = [];
 const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false};
 function setState(ns) {
   const { ctx:prevCtx, step:prevStep } = select(state.currentCtx, state.currentStep);
+  const prevRewrite = state.currentRewrite;
   Object.assign(state, ns);
   // update element styles if needed
   const { ctx, step } = select(state.currentCtx, state.currentStep);
   toggleCls(prevCtx, ctx, "expanded", state.expandSteps);
   if (ctx?.id !== prevCtx?.id) {
-    saveToHistory({ currentCtx:deselect(prevCtx).ctx, currentRewrite:0, currentStep:0, expandSteps:false });
+    saveToHistory({ currentCtx:deselect(prevCtx).ctx, currentStep:deselect(prevStep).step || 0, currentRewrite:prevRewrite, expandSteps:true });
     toggleCls(prevCtx, ctx, "active");
   }
   if (ctx?.id !== prevCtx?.id || step?.id !== prevStep?.id) {
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 5af45132ab..8e2bead195 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -285,7 +285,7 @@ class Handler(BaseHTTPRequestHandler):
       if url.path == "/render": ret, content_type = json.dumps(get_render(get_int(query, "ctx"), query["fmt"][0])).encode(), "application/json"
       else:
         try: return self.stream_json(get_full_rewrite(trace.rewrites[i:=get_int(query, "ctx")][get_int(query, "idx")], i))
-        except KeyError: status_code = 404
+        except (KeyError, IndexError): status_code = 404
     elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
     elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
     else: status_code = 404