diff --git a/opencl/test/unit_test/helpers/timestamp_packet_tests.h b/opencl/test/unit_test/helpers/timestamp_packet_tests.h
index 391f468975..b2d168dda1 100644
--- a/opencl/test/unit_test/helpers/timestamp_packet_tests.h
+++ b/opencl/test/unit_test/helpers/timestamp_packet_tests.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -34,6 +34,7 @@ struct TimestampPacketTests : public ::testing::Test {
             executionEnvironment->rootDeviceEnvironments[i]->setHwInfoAndInitHelpers(defaultHwInfo.get());
             executionEnvironment->rootDeviceEnvironments[i]->initGmm();
         }
+        executionEnvironment->calculateMaxOsContextCount();
         device = std::make_unique<MockClDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
         context = new MockContext(device.get());
         kernel = std::make_unique<MockKernelWithInternals>(*device, context);
diff --git a/shared/source/memory_manager/graphics_allocation.cpp b/shared/source/memory_manager/graphics_allocation.cpp
index 40c821311b..c83837f4db 100644
--- a/shared/source/memory_manager/graphics_allocation.cpp
+++ b/shared/source/memory_manager/graphics_allocation.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -55,6 +55,7 @@ GraphicsAllocation::GraphicsAllocation(uint32_t rootDeviceIndex, size_t numGmms,
 GraphicsAllocation::~GraphicsAllocation() = default;
 
 void GraphicsAllocation::updateTaskCount(TaskCountType newTaskCount, uint32_t contextId) {
+    UNRECOVERABLE_IF(contextId >= usageInfos.size());
     if (usageInfos[contextId].taskCount == objectNotUsed) {
         registeredContextsNum++;
     }
diff --git a/shared/source/memory_manager/graphics_allocation.h b/shared/source/memory_manager/graphics_allocation.h
index 96efe2daef..d3375d3096 100644
--- a/shared/source/memory_manager/graphics_allocation.h
+++ b/shared/source/memory_manager/graphics_allocation.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -160,7 +160,13 @@ class GraphicsAllocation : public IDNode<GraphicsAllocation> {
     bool isUsedByManyOsContexts() const { return registeredContextsNum > 1u; }
     bool isUsedByOsContext(uint32_t contextId) const { return objectNotUsed != getTaskCount(contextId); }
     MOCKABLE_VIRTUAL void updateTaskCount(TaskCountType newTaskCount, uint32_t contextId);
-    MOCKABLE_VIRTUAL TaskCountType getTaskCount(uint32_t contextId) const { return usageInfos[contextId].taskCount; }
+    MOCKABLE_VIRTUAL TaskCountType getTaskCount(uint32_t contextId) const {
+        if (contextId >= usageInfos.size()) {
+            return objectNotUsed;
+        }
+
+        return usageInfos[contextId].taskCount;
+    }
     void releaseUsageInOsContext(uint32_t contextId) { updateTaskCount(objectNotUsed, contextId); }
     uint32_t getInspectionId(uint32_t contextId) const { return usageInfos[contextId].inspectionId; }
     void setInspectionId(uint32_t newInspectionId, uint32_t contextId) { usageInfos[contextId].inspectionId = newInspectionId; }
diff --git a/shared/test/unit_test/device/neo_device_tests.cpp b/shared/test/unit_test/device/neo_device_tests.cpp
index 148749d06a..d074e50ae9 100644
--- a/shared/test/unit_test/device/neo_device_tests.cpp
+++ b/shared/test/unit_test/device/neo_device_tests.cpp
@@ -1080,6 +1080,8 @@ HWCMDTEST_F(IGFX_XE_HP_CORE, DeviceTests, givenCCSEngineAndContextGroupSizeEnabl
 
     UltDeviceFactory deviceFactory{1, 0, executionEnvironment};
 
+    MemoryManager::maxOsContextCount++;
+
     deviceFactory.rootDevices[0]->createEngine(0, {aub_stream::EngineType::ENGINE_CCS, EngineUsage::regular});
 
     auto defaultEngine = deviceFactory.rootDevices[0]->getDefaultEngine();
diff --git a/shared/test/unit_test/memory_manager/graphics_allocation_tests.cpp b/shared/test/unit_test/memory_manager/graphics_allocation_tests.cpp
index bd7da5cd94..662c9498b3 100644
--- a/shared/test/unit_test/memory_manager/graphics_allocation_tests.cpp
+++ b/shared/test/unit_test/memory_manager/graphics_allocation_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2023 Intel Corporation
+ * Copyright (C) 2018-2024 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -59,6 +59,20 @@ TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatedTaskCountThenOnly
         EXPECT_EQ(MockGraphicsAllocation::objectNotUsed, graphicsAllocation.getTaskCount(i));
     }
 }
+
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenGettingTaskCountForInvalidContextIdThenObjectNotUsedIsReturned) {
+    MockGraphicsAllocation graphicsAllocation;
+    for (auto i = 0u; i < MemoryManager::maxOsContextCount; i++) {
+        graphicsAllocation.updateTaskCount(1u, i);
+    }
+    EXPECT_EQ(MockGraphicsAllocation::objectNotUsed, graphicsAllocation.getTaskCount(MemoryManager::maxOsContextCount));
+}
+
+TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatingTaskCountForInvalidContextIdThenAbortExecution) {
+    MockGraphicsAllocation graphicsAllocation;
+    EXPECT_THROW(graphicsAllocation.updateTaskCount(1u, MemoryManager::maxOsContextCount), std::runtime_error);
+}
+
 TEST(GraphicsAllocationTest, givenGraphicsAllocationWhenUpdatedTaskCountToobjectNotUsedValueThenUnregisterContext) {
     MockGraphicsAllocation graphicsAllocation;
     EXPECT_FALSE(graphicsAllocation.isUsed());
@@ -545,4 +559,4 @@ TEST(GraphicsAllocationTest, givenGraphicsAllocationsWithFragmentsWhenCallingFor
         auto residencyData = graphicsAllocation.fragmentsStorage.fragmentStorageData[allocationId].residency;
         EXPECT_EQ(residencyData->getFenceValueForContextId(contextId), newFenceValue);
     }
-}
\ No newline at end of file
+}
diff --git a/shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp
index e1c5d4494b..655d0694a4 100644
--- a/shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_memory_manager_localmem_prelim_tests.cpp
@@ -683,6 +683,7 @@ TEST_F(DrmMemoryManagerLocalMemoryPrelimTest, givenUseKmdMigrationSetWhenCreateS
     mock->ioctlCallsCount = 0;
     mock->setBindAvailable();
 
+    executionEnvironment->calculateMaxOsContextCount();
     SVMAllocsManager unifiedMemoryManager(memoryManager, false);
 
     SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::sharedUnifiedMemory, 1, rootDeviceIndices, deviceBitfields);