diff --git a/core/libs/prx/libSceAgcDriver/CMakeLists.txt b/core/libs/prx/libSceAgcDriver/CMakeLists.txt
index 9d27611288..2a8d0e5d55 100644
--- a/core/libs/prx/libSceAgcDriver/CMakeLists.txt
+++ b/core/libs/prx/libSceAgcDriver/CMakeLists.txt
@@ -2036,6 +2036,17 @@ add_test_executable(agc_driver_graphics_tests EXCLUDE_FROM_ALL
 )
 
 if(BUILD_TESTING)
+    add_test_executable(agc_unit_flags_allocation_failure_tests
+        tests/UnitFlagsAllocationFailure.cpp tests/AllocationFailure.cpp)
+    target_include_directories(agc_unit_flags_allocation_failure_tests PRIVATE
+        $<TARGET_PROPERTY:libSceAgcDriver,INCLUDE_DIRECTORIES>)
+    if(MINGW)
+        target_compile_options(agc_unit_flags_allocation_failure_tests PRIVATE -fasynchronous-unwind-tables)
+    endif()
+    configure_windows_unwind(agc_unit_flags_allocation_failure_tests)
+    add_test(NAME agc_unit_flags_allocation_failure_empty COMMAND agc_unit_flags_allocation_failure_tests empty)
+    add_test(NAME agc_unit_flags_allocation_failure_populated COMMAND agc_unit_flags_allocation_failure_tests populated)
+    set_tests_properties(agc_unit_flags_allocation_failure_empty agc_unit_flags_allocation_failure_populated PROPERTIES TIMEOUT 20)
     add_test_executable(agc_descriptor_cache_failure_tests
         tests/DescriptorCacheFailure.cpp
         Graphics/src/ShaderResources.cpp)
diff --git a/core/libs/prx/libSceAgcDriver/Graphics/include/Texture.hpp b/core/libs/prx/libSceAgcDriver/Graphics/include/Texture.hpp
index 0f20bbdef2..bb90bae546 100644
--- a/core/libs/prx/libSceAgcDriver/Graphics/include/Texture.hpp
+++ b/core/libs/prx/libSceAgcDriver/Graphics/include/Texture.hpp
@@ -186,9 +186,9 @@ public:
     };
 
     void assign(std::size_t size, bool flag) {
-        count = size;
         words.assign((size + WordBits - 1u) / WordBits, flag ? ~std::uint64_t{0} : 0u);
         if (flag && size % WordBits != 0u) words.back() &= (std::uint64_t{1} << (size % WordBits)) - 1u;
+        count = size;
     }
     void clear() {
         words.clear();
diff --git a/core/libs/prx/libSceAgcDriver/tests/UnitFlagsAllocationFailure.cpp b/core/libs/prx/libSceAgcDriver/tests/UnitFlagsAllocationFailure.cpp
new file mode 100644
index 0000000000..3d8001b733
--- /dev/null
+++ b/core/libs/prx/libSceAgcDriver/tests/UnitFlagsAllocationFailure.cpp
@@ -0,0 +1,67 @@
+#include "AllocationFailure.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/Texture.hpp"
+#include <cstdio>
+#include <new>
+#include <stdexcept>
+#include <string_view>
+
+namespace {
+
+void require(bool condition, const char* message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+void failAssign(AgcDriver::Graphics::UnitFlags& flags, std::size_t size) {
+    bool failed = false;
+    AgcDriver::Tests::SetAllocationFailureCountdown(0);
+    try {
+        flags.assign(size, true);
+    } catch (const std::bad_alloc&) {
+        failed = true;
+    } catch (...) {
+        AgcDriver::Tests::SetAllocationFailureCountdown(-1);
+        throw;
+    }
+    AgcDriver::Tests::SetAllocationFailureCountdown(-1);
+    require(failed, "unit flags assignment did not exercise allocation failure");
+}
+
+void emptySet() {
+    AgcDriver::Graphics::UnitFlags flags;
+    failAssign(flags, 65);
+    require(flags.empty() && flags.size() == 0 && !flags.Any(), "failed first allocation changed the unit flags size");
+    flags.assign(65, true);
+    require(flags.size() == 65 && flags.Count() == 65 && flags.AllIn(0, 65), "unit flags did not recover after a failed first allocation");
+    flags.SetRange(64, 1, false);
+    require(flags.Count() == 64 && !flags.AnyIn(64, 1), "recovered unit flags lost range access");
+}
+
+void populatedSet() {
+    AgcDriver::Graphics::UnitFlags flags;
+    flags.assign(65, false);
+    flags[3] = true;
+    flags[64] = true;
+    failAssign(flags, 129);
+    require(flags.size() == 65 && flags.Count() == 2 && flags[3] && flags[64], "failed growth changed the existing unit flags");
+    flags.assign(130, true);
+    require(flags.size() == 130 && flags.Count() == 130 && flags.AllIn(0, 130), "unit flags did not recover after failed growth");
+    flags.SetRange(0, 130, false);
+    require(!flags.Any(), "recovered unit flags retained bits after clearing");
+}
+
+}
+
+int main(int argc, char** argv) {
+    try {
+        require(argc == 2, "expected empty or populated case");
+        const std::string_view mode(argv[1]);
+        if (mode == "empty") emptySet();
+        else if (mode == "populated") populatedSet();
+        else throw std::runtime_error("unknown unit flags allocation-failure case");
+        return 0;
+    } catch (const std::exception& error) {
+        AgcDriver::Tests::SetAllocationFailureCountdown(-1);
+        std::fprintf(stderr, "%s\n", error.what());
+        return 1;
+    }
+}