Preserve yield state across caught callback errors

VM/include/lyieldable.h

@@ -142,6 +142,19 @@ class SlotManager
     // frame instead, so there's nothing to restore.
     void flushForYield(bool own_frame);
 
+    // Saves the state ahead of a call to the function at the top of the stack under `nargs`
+    // arguments that wants `nresults` results. The frame top it records is the function's slot plus
+    // the results: that is where a callee that returns after an error caught inside it leaves the
+    // frame, instead of the frame limit, which the continuation would otherwise find above the
+    // results.
+    void flushForCall(int nargs, int nresults);
+
+    // Undoes the "yielding" mark flushForYield() put on this manager and its parents, for a
+    // call that saved the state ahead of time and then returned normally: the frames are not
+    // being unwound, so their destructors must neither write the buffer nor leave a parent
+    // marked as having a child.
+    void clearYielding();
+
     template<typename T>
     PrimitiveSlot<T> allocSlot(T* storage);
 };

VM/include/lyieldablemacros.h

@@ -136,6 +136,30 @@ using Luau::YieldGuard;
     _yieldable_label_##phase_name: ;                                                \
     } while (0)
 
+// Like YIELD_CALL, for a function whose callee runs script code that may catch an error with
+// pcall. Such a catch unwinds this frame (and the frames of any chained helpers) without a yield,
+// and the continuation that finishes the function afterwards reads the state saved at the call,
+// so the state is saved before the call and the "yielding" mark is cleared again when the callee
+// returns normally.
+#define YIELD_CALL_SAVING_STATE(L, nargs, nresults, phase_name)                 \
+    do                                                                          \
+    {                                                                           \
+        (_yieldable_phase) = Phase::phase_name;                                 \
+        _yieldable_slots.flushForCall(nargs, nresults);                         \
+        lua_call(L, nargs, nresults);                                           \
+        if ((L)->status != LUA_OK)                                              \
+        {                                                                       \
+            if ((L)->status == LUA_YIELD || (L)->status == LUA_BREAK)           \
+            {                                                                   \
+                _yieldable_slots.flushForYield(false);                          \
+                return _yieldable_on_yield();                                   \
+            }                                                                   \
+            return _yieldable_on_error();                                       \
+        }                                                                       \
+        _yieldable_slots.clearYielding();                                       \
+    _yieldable_label_##phase_name: ;                                            \
+    } while (0)
+
 // Calls a helper function that has its own chained SlotManager and dispatch
 // region. The goto label is BEFORE the call so the helper is re-called on
 // resume. No flushForYield here — the helper already flushed the entire chain.

VM/src/cjson/lua_cjson.cpp

@@ -503,7 +503,7 @@ static void json_append_array(lua_State* l, SlotManager& parent_slots,
                 lua_pushvalue(l, -2);                        // self
                 lua_pushvalue(l, (int)EncodeStack::CTX);     // ctx
                 YIELD_CHECK(l, TOJSON_CHECK, LUA_INTERRUPT_LLLIB);
-                YIELD_CALL(l, 2, 1, TOJSON_CALL);
+                YIELD_CALL_SAVING_STATE(l, 2, 1, TOJSON_CALL);
                 // Stack: table, value, resolved
                 lua_remove(l, -2);  // remove original value
                 // Stack: table, resolved
@@ -514,7 +514,7 @@ static void json_append_array(lua_State* l, SlotManager& parent_slots,
             lua_pushvalue(l, -3);      // value (possibly __tojson-resolved)
             lua_pushvalue(l, -5);      // parent (the array table)
             YIELD_CHECK(l, REPLACER_CHECK, LUA_INTERRUPT_LLLIB);
-            YIELD_CALL(l, 3, 1, REPLACER_CALL);
+            YIELD_CALL_SAVING_STATE(l, 3, 1, REPLACER_CALL);
             // Stack: table, value, result
             if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
                 lua_pop(l, 2);  // pop result + value -> table
@@ -846,7 +846,7 @@ static void json_append_object(lua_State* l, SlotManager& parent_slots,
                 lua_pushvalue(l, -2);                        // self
                 lua_pushvalue(l, (int)EncodeStack::CTX);     // ctx
                 YIELD_CHECK(l, TOJSON_CHECK, LUA_INTERRUPT_LLLIB);
-                YIELD_CALL(l, 2, 1, TOJSON_CALL);
+                YIELD_CALL_SAVING_STATE(l, 2, 1, TOJSON_CALL);
                 // Stack: table, key, value, resolved
                 lua_remove(l, -2);  // remove original value
                 // Stack: table, key, resolved
@@ -857,7 +857,7 @@ static void json_append_object(lua_State* l, SlotManager& parent_slots,
             lua_pushvalue(l, -3);  // value (possibly __tojson-resolved)
             lua_pushvalue(l, -6);  // parent (the object table)
             YIELD_CHECK(l, REPLACER_CHECK, LUA_INTERRUPT_LLLIB);
-            YIELD_CALL(l, 3, 1, REPLACER_CALL);
+            YIELD_CALL_SAVING_STATE(l, 3, 1, REPLACER_CALL);
             // Stack: table, key, value, result
             if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
                 lua_pop(l, 2);  // pop result + value -> table, key
@@ -1064,7 +1064,7 @@ static int json_append_data(lua_State* l, SlotManager& parent_slots,
                 lua_pushvalue(l, -2);                           // self
                 lua_pushvalue(l, (int)EncodeStack::CTX);        // ctx table
                 YIELD_CHECK(l, TOJSON_CHECK, LUA_INTERRUPT_LLLIB);
-                YIELD_CALL(l, 2, 1, TOJSON_CALL);
+                YIELD_CALL_SAVING_STATE(l, 2, 1, TOJSON_CALL);
                 // Stack: ..., original_table, tojson_result
                 if (lua_istable(l, -1)) {
                     // Table result: replace original, fall through to shape handling
@@ -1301,7 +1301,7 @@ static int json_encode_common(lua_State* l, bool is_init, uint8_t abi_version, b
         lua_pushvalue(l, (int)EncodeStack::VALUE);     // self
         lua_pushvalue(l, (int)EncodeStack::CTX);       // ctx
         YIELD_CHECK(l, ROOT_TOJSON_CHECK, LUA_INTERRUPT_LLLIB);
-        YIELD_CALL(l, 2, 1, ROOT_TOJSON_CALL);
+        YIELD_CALL_SAVING_STATE(l, 2, 1, ROOT_TOJSON_CALL);
         // Stack: [..., resolved]
         lua_replace(l, (int)EncodeStack::VALUE);
     }
@@ -1313,7 +1313,7 @@ static int json_encode_common(lua_State* l, bool is_init, uint8_t abi_version, b
         lua_pushvalue(l, (int)EncodeStack::VALUE);     // value (possibly __tojson-resolved)
         lua_pushnil(l);                                // parent = nil (root)
         YIELD_CHECK(l, ROOT_REPLACER_CHECK, LUA_INTERRUPT_LLLIB);
-        YIELD_CALL(l, 3, 1, ROOT_REPLACER_CALL);
+        YIELD_CALL_SAVING_STATE(l, 3, 1, ROOT_REPLACER_CALL);
         // Stack: [..., result]
         if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
             lua_pop(l, 1);
@@ -2093,7 +2093,7 @@ static void json_parse_object_context(lua_State* l, SlotManager& parent_slots, j
             lua_pushvalue(l, -6);  // parent (the object table)
             lua_pushvalue(l, (int)DecodeStack::CTX);
             YIELD_CHECK(l, REVIVER_CHECK, LUA_INTERRUPT_LLLIB);
-            YIELD_CALL(l, 4, 1, REVIVER_CALL);
+            YIELD_CALL_SAVING_STATE(l, 4, 1, REVIVER_CALL);
             // Stack: [..., table, key, value, result]
             if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
                 // Omit this key/value pair entirely
@@ -2216,7 +2216,7 @@ static void json_parse_array_context(lua_State* l, SlotManager& parent_slots, js
             lua_pushvalue(l, -5);      // parent (the array table)
             lua_pushvalue(l, (int)DecodeStack::CTX);
             YIELD_CHECK(l, REVIVER_CHECK, LUA_INTERRUPT_LLLIB);
-            YIELD_CALL(l, 4, 1, REVIVER_CALL);
+            YIELD_CALL_SAVING_STATE(l, 4, 1, REVIVER_CALL);
             // Stack: [..., table, value, result]
             if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
                 // Omit this element - don't insert, don't increment insert_idx
@@ -2472,7 +2472,7 @@ static int json_decode_common(lua_State* l, bool is_init, uint8_t abi_version, b
         lua_pushnil(l);        // parent = nil (root)
         lua_pushvalue(l, (int)DecodeStack::CTX);
         YIELD_CHECK(l, ROOT_REVIVER_CHECK, LUA_INTERRUPT_LLLIB);
-        YIELD_CALL(l, 4, 1, ROOT_REVIVER_CALL);
+        YIELD_CALL_SAVING_STATE(l, 4, 1, ROOT_REVIVER_CALL);
         // Stack: [..., decoded_value, result]
         if (lua_tolightuserdatatagged(l, -1, LU_TAG_JSON_INTERNAL) == JSON_REMOVE) {
             lua_pop(l, 2);

VM/src/lllevents.cpp

@@ -673,7 +673,7 @@ DEFINE_YIELDABLE(llevents_handle_event, 0)
         }
 
         // handler function and arguments are in order on the stack, do the call.
-        YIELD_CALL(L, nargs, 0, CALL_HANDLER);
+        YIELD_CALL_SAVING_STATE(L, nargs, 0, CALL_HANDLER);
     }
 
     // All handlers completed - mark DetectedEvent wrappers as invalid

VM/src/llltimers.cpp

@@ -672,7 +672,7 @@ DEFINE_YIELDABLE(lltimers_tick, 0)
                 lua_pushnumber(L, interval);
             }
         }
-        YIELD_CALL(L, 2, 0, CALL_HANDLER);
+        YIELD_CALL_SAVING_STATE(L, 2, 0, CALL_HANDLER);
 
         // Nil out the timer reference, we don't need it anymore
         lua_pushnil(L);

VM/src/ltablib.cpp

@@ -506,7 +506,7 @@ static bool sort_cmp_may_call(LuaTable* t, int i, int j)
             setobj2s(L, L->top + 2, &t->array[j_idx]);                                          \
             L->top += 3;                                                                         \
         }                                                                                        \
-        YIELD_CALL(L, 2, 1, phase_name);                                                         \
+        YIELD_CALL_SAVING_STATE(L, 2, 1, phase_name);                                                         \
         cmp_var = lua_toboolean(L, -1) != 0;                                                     \
         lua_pop(L, 1);                                                                           \
         if (t->sizearray != saved_sa)                                                            \

VM/src/lyieldable.cpp

@@ -80,6 +80,23 @@ SlotManager::SlotManager(lua_State* L, bool is_init, uint8_t abi_version)
         luaL_checkstack(L, needed, "yieldable frame");
 }
 
+void SlotManager::flushForCall(int nargs, int nresults)
+{
+    LUAU_ASSERT(nresults >= 0);
+    flushForYield(false);
+
+    ptrdiff_t frame_size = (L->top - nargs - 1 + nresults) - (L->stack + bufferStackOffset);
+    LUAU_ASSERT(frame_size >= 1 && frame_size < UINT16_MAX);
+    uint16_t top = (uint16_t)frame_size;
+    memcpy(bufferData + Luau::kYieldStateHeader, &top, sizeof(uint16_t));
+}
+
+void SlotManager::clearYielding()
+{
+    for (SlotManager* mgr = this; mgr; mgr = mgr->parent)
+        mgr->yielding = false;
+}
+
 // Creates a yield state userdata at position 1 sized for the entire chain.
 // Sets bufferData and yielding on all managers so slot destructors
 // write on unwind.

VM/src/lyieldstrlib.cpp

@@ -1408,7 +1408,7 @@ DEFINE_YIELDABLE_EXTERN(yieldable_str_gsub, 0)
                 }
             }
             if (repl_type == LUA_TFUNCTION)
-                YIELD_CALL(L, call_nargs, 1, REPL_CALL);
+                YIELD_CALL_SAVING_STATE(L, call_nargs, 1, REPL_CALL);
             // Function/table result: nil or false -> keep original match text
             if (!lua_toboolean(L, -1))
             {

tests/SLExecutor.test.cpp

@@ -3462,3 +3462,138 @@ TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor serialize mid-handler roundtrip")
 }
 
 TEST_SUITE_END();
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor a handler that catches an error with pcall finishes and the next event runs")
+{
+    TestScript ts(R"(
+        function LLEvents.moving_start()
+            local ok = pcall(function() error("boom") end)
+            print(`caught {ok}`)
+        end
+    )");
+    ts.start();
+    Script& exec = ts.exec;
+
+    RunResult first;
+    {
+        RunWindow window(exec, 1.0);
+        first = exec.callEventHandler(LSLEvent::MovingStart, nullptr);
+    }
+    CHECK(first.status == HandlerRunStatus::Ok);
+    CHECK(exec.getFaultKind() == FaultKind::None);
+
+    RunResult second;
+    {
+        RunWindow window(exec, 1.0);
+        second = exec.callEventHandler(LSLEvent::MovingStart, nullptr);
+    }
+    CHECK(second.status == HandlerRunStatus::Ok);
+    CHECK(exec.getFaultKind() == FaultKind::None);
+    checkCapture(ts.host.printed, {"caught false", "caught false"});
+}
+
+// An error caught with pcall inside the callee of a yieldable library function unwinds the
+// function's frame without a yield; the function must still finish when the callee returns.
+// Each case runs the same event twice and expects no fault and the same output both times.
+static void checkCaughtErrorInCallee(const char* body, const char* expected, bool catches = true)
+{
+    const std::string source = std::string(catches ? "local function swallow() pcall(function() error('boom') end) end\n"
+                                                   : "local function swallow() end\n") +
+                               "function LLEvents.moving_start()\n" + body + "\nend\n";
+    TestScript ts(source.c_str());
+    ts.start();
+    Script& exec = ts.exec;
+
+    for (int i = 0; i < 2; ++i)
+    {
+        RunResult result;
+        {
+            RunWindow window(exec, 1.0);
+            result = exec.callEventHandler(LSLEvent::MovingStart, nullptr);
+        }
+        CHECK(result.status == HandlerRunStatus::Ok);
+        INFO("fault: ", exec.getFaultString());
+        CHECK(exec.getFaultKind() == FaultKind::None);
+    }
+    checkCapture(ts.host.printed, {expected, expected});
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor string.gsub with a function that catches an error with pcall finishes")
+{
+    checkCaughtErrorInCallee(R"(
+        local s = string.gsub("abc", "%w", function(c) swallow() return c:upper() end)
+        print(s)
+    )", "ABC");
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor table.sort with a comparator that catches an error with pcall finishes")
+{
+    checkCaughtErrorInCallee(R"(
+        local t = {3, 1, 2}
+        table.sort(t, function(a, b) swallow() return a < b end)
+        print(table.concat(t, ","))
+    )", "1,2,3");
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor lljson.encode with a replacer that catches an error with pcall finishes")
+{
+    checkCaughtErrorInCallee(R"(
+        print(lljson.encode({1, 2}, {replacer = function(key, value) swallow() return value end}))
+    )", "[1,2]");
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor lljson.encode with a __tojson that catches an error with pcall finishes")
+{
+    checkCaughtErrorInCallee(R"(
+        local obj = setmetatable({}, {__tojson = function(self) swallow() return "ok" end})
+        print(lljson.encode({x = obj}))
+    )", "{\"x\":\"ok\"}");
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor lljson.decode with a reviver that catches an error with pcall finishes")
+{
+    checkCaughtErrorInCallee(R"(
+        local t = lljson.decode('{"a":1}', function(key, value) swallow() return value end)
+        print(t.a)
+    )", "1");
+}
+
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor a timer handler that catches an error with pcall finishes and the next tick runs")
+{
+    TestScript ts(R"(
+        LLTimers:every(1.0, function()
+            pcall(function() error("boom") end)
+            print("tick")
+        end)
+    )");
+    ts.loadDefaultState();
+    resume(ts.exec);
+
+    for (double when : {1.1, 2.5})
+    {
+        ts.host.script_clock = when;
+        RunResult result = dispatchRaw(ts.exec, LSLEvent::Timer);
+        if (result.status == HandlerRunStatus::Preempted)
+            resumeToCompletion(ts.exec, 1.0);
+        else
+            CHECK(result.status == HandlerRunStatus::Ok);
+        INFO("fault: ", ts.exec.getFaultString());
+        CHECK(ts.exec.getFaultKind() == FaultKind::None);
+    }
+    checkCapture(ts.host.printed, {"tick", "tick"});
+}
+
+// The control for the cases above: the same callbacks with nothing caught finish, so those cases
+// fail only because of the caught error.
+TEST_CASE_FIXTURE(SLuaFixture, "SLExecutor yieldable library callbacks that catch nothing finish")
+{
+    checkCaughtErrorInCallee(R"(
+        local s = string.gsub("abc", "%w", function(c) swallow() return c:upper() end)
+        local t = {3, 1, 2}
+        table.sort(t, function(a, b) swallow() return a < b end)
+        local e = lljson.encode({1, 2}, {replacer = function(key, value) swallow() return value end})
+        local o = setmetatable({}, {__tojson = function(self) swallow() return "ok" end})
+        local d = lljson.decode('{"a":1}', function(key, value) swallow() return value end)
+        print(s .. "|" .. table.concat(t, ",") .. "|" .. e .. "|" .. lljson.encode({x = o}) .. "|" .. d.a)
+    )", "ABC|1,2,3|[1,2]|{\"x\":\"ok\"}|1", false);
+}