fix: preserve nested JS return expressions

src/browser_harness/helpers.py

@@ -117,43 +117,13 @@ def _runtime_evaluate(expression, session_id=None, await_promise=False):
     return _runtime_value(r, expression)
 
 
-def _has_return_statement(expression):
-    i = 0
-    n = len(expression)
-    state = "code"
-    quote = ""
-    while i < n:
-        ch = expression[i]
-        nxt = expression[i + 1] if i + 1 < n else ""
-        if state == "code":
-            if ch in ("'", '"', "`"):
-                state = "string"; quote = ch; i += 1; continue
-            if ch == "/" and nxt == "/":
-                state = "line_comment"; i += 2; continue
-            if ch == "/" and nxt == "*":
-                state = "block_comment"; i += 2; continue
-            if expression.startswith("return", i):
-                before = expression[i - 1] if i > 0 else ""
-                after = expression[i + 6] if i + 6 < n else ""
-                if not (before == "_" or before.isalnum()) and not (after == "_" or after.isalnum()):
-                    return True
-            i += 1; continue
-        if state == "line_comment":
-            if ch == "\n":
-                state = "code"
-            i += 1; continue
-        if state == "block_comment":
-            if ch == "*" and nxt == "/":
-                state = "code"; i += 2; continue
-            i += 1; continue
-        if state == "string":
-            if ch == "\\":
-                i += 2; continue
-            if ch == quote:
-                state = "code"; quote = ""
-            i += 1; continue
-    return False
+def _wrap_js_return_statement(expression):
+    return f"(function(){{{expression}}})()"
+
 
+def _is_illegal_return_statement(error):
+    message = str(error)
+    return "SyntaxError" in message and "Illegal return statement" in message
 
 # --- navigation / page ---
 def goto_url(url):
@@ -426,13 +396,18 @@ def wait_for_network_idle(timeout=10.0, idle_ms=500):
 def js(expression, target_id=None):
     """Run JS in the attached tab (default) or inside an iframe target (via iframe_target()).
 
-    Expressions with top-level `return` are automatically wrapped in an IIFE, so both
-    `document.title` and `const x = 1; return x` are valid inputs.
+    The browser accepts expression snippets directly. If Chrome reports a
+    top-level `return` syntax error, retry as an IIFE so snippets like
+    `const x = 1; return x` still work without mistaking nested function returns
+    for top-level returns.
     """
     sid = cdp("Target.attachToTarget", targetId=target_id, flatten=True)["sessionId"] if target_id else None
-    if _has_return_statement(expression) and not expression.strip().startswith("("):
-        expression = f"(function(){{{expression}}})()"
-    return _runtime_evaluate(expression, session_id=sid, await_promise=True)
+    try:
+        return _runtime_evaluate(expression, session_id=sid, await_promise=True)
+    except RuntimeError as e:
+        if expression.strip().startswith("(") or not _is_illegal_return_statement(e):
+            raise
+        return _runtime_evaluate(_wrap_js_return_statement(expression), session_id=sid, await_promise=True)
 
 
 _KC = {"Enter": 13, "Tab": 9, "Escape": 27, "Backspace": 8, " ": 32, "ArrowLeft": 37, "ArrowUp": 38, "ArrowRight": 39, "ArrowDown": 40}

tests/integration/test_js.py

@@ -19,18 +19,45 @@ def _evaluated_expression(captured):
     return next(kw["expression"] for m, kw in captured if m == "Runtime.evaluate")
 
 
+def _evaluated_expressions(captured):
+    return [kw["expression"] for m, kw in captured if m == "Runtime.evaluate"]
+
+
+def _illegal_return_response():
+    return {
+        "result": {
+            "type": "object",
+            "subtype": "error",
+            "description": "SyntaxError: Illegal return statement",
+        },
+        "exceptionDetails": {
+            "text": "Uncaught SyntaxError: Illegal return statement",
+            "lineNumber": 0,
+            "columnNumber": 0,
+        },
+    }
+
+
 def test_simple_expression_passes_through():
     fake_cdp, captured = _capture_cdp()
     with patch("browser_harness.helpers.cdp", side_effect=fake_cdp):
         helpers.js("document.title")
     assert _evaluated_expression(captured) == "document.title"
 
 
-def test_return_statement_gets_wrapped():
-    fake_cdp, captured = _capture_cdp()
+def test_return_statement_gets_wrapped_after_illegal_return():
+    captured = []
+    expr = "const x = 1; return x"
+    responses = [_illegal_return_response(), {"result": {"value": 1}}]
+
+    def fake_cdp(method, **kwargs):
+        captured.append((method, kwargs))
+        return responses.pop(0)
+
     with patch("browser_harness.helpers.cdp", side_effect=fake_cdp):
-        helpers.js("const x = 1; return x")
-    assert _evaluated_expression(captured) == "(function(){const x = 1; return x})()"
+        assert helpers.js(expr) == 1
+
+    assert _evaluated_expressions(captured) == [expr, "(function(){const x = 1; return x})()"]
 
 
 def test_iife_with_internal_return_is_not_double_wrapped():
@@ -40,6 +67,14 @@ def test_iife_with_internal_return_is_not_double_wrapped():
     assert _evaluated_expression(captured) == "(function(){ return document.title; })()"
 
 
+def test_nested_arrow_return_is_not_wrapped():
+    fake_cdp, captured = _capture_cdp()
+    expr = "const f = () => { return 1 }; f()"
+    with patch("browser_harness.helpers.cdp", side_effect=fake_cdp):
+        helpers.js(expr)
+    assert _evaluated_expression(captured) == expr
+
+
 def test_js_raises_on_syntax_error_exception_details():
     def fake_cdp(method, **kwargs):
         return {
@@ -112,11 +147,18 @@ def test_return_word_inside_comment_does_not_trigger_wrapping():
 
 
 @pytest.mark.parametrize("expr", ["return\t1", "return\n1"])
-def test_top_level_return_with_whitespace_gets_wrapped(expr):
-    fake_cdp, captured = _capture_cdp()
+def test_top_level_return_with_whitespace_gets_wrapped_after_illegal_return(expr):
+    captured = []
+    responses = [_illegal_return_response(), {"result": {"value": 1}}]
+
+    def fake_cdp(method, **kwargs):
+        captured.append((method, kwargs))
+        return responses.pop(0)
+
     with patch("browser_harness.helpers.cdp", side_effect=fake_cdp):
-        helpers.js(expr)
-    assert _evaluated_expression(captured) == f"(function(){{{expr}}})()"
+        assert helpers.js(expr) == 1
+
+    assert _evaluated_expressions(captured) == [expr, f"(function(){{{expr}}})()"]
 
 
 @pytest.mark.parametrize(