fix bug

Signed-off-by: xhz-sz <xiehuandong@h-partners.com>

AI[91%] Human Fixed[0%] Human[9%] AI Adopted[100%]
Co-authored-by: codex (gpt-5.6-sol) <ai@local>
Co-authored-by: codex (unknown) <ai@local>
Co-authored-by: opencode (glm-5.2) <ai@local>

frameworks/js/napi/ability_manager/js_preload_ui_extension_callback_client.cpp

@@ -90,6 +90,10 @@ void JsPreloadUIExtensionCallbackClient::ProcessOnDestroyDone(int32_t extensionA
 
 void JsPreloadUIExtensionCallbackClient::CallJsPreloadedUIExtensionAbility(int32_t preloadId)
 {
+    if (env_ == nullptr) {
+        TAG_LOGE(AAFwkTag::ABILITYMGR, "null env_");
+        return;
+    }
     HandleScope handleScope(env_);
     if (callbackRef_ == nullptr) {
         TAG_LOGE(AAFwkTag::ABILITYMGR, "null callbackRef_");

frameworks/native/ability/native/ui_extension_base/js_ui_extension_context.cpp

@@ -1133,6 +1133,26 @@ bool JsUIExtensionContext::CheckConnectAlreadyExist(napi_env env, AAFwk::Want& w
     return true;
 }
 
+static sptr<JSUIServiceUIExtConnection> CreateUIServiceExtConnection(napi_env env, AAFwk::Want& want)
+{
+    sptr<JSUIServiceUIExtConnection> connection = sptr<JSUIServiceUIExtConnection>::MakeSptr(env);
+    if (connection == nullptr) {
+        TAG_LOGE(AAFwkTag::UISERVC_EXT, "null connection");
+        ThrowError(env, static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER),
+            GetInnerErrorMsg(AbilityInnerErrorMsg::MEMORY_ALLOC_FAILED));
+        return nullptr;
+    }
+    sptr<UIExtensionServiceHostStubImpl> stub = connection->GetServiceHostStub();
+    if (stub == nullptr) {
+        TAG_LOGE(AAFwkTag::UISERVC_EXT, "null service host stub");
+        ThrowError(env, static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER),
+            GetInnerErrorMsg(AbilityInnerErrorMsg::MEMORY_ALLOC_FAILED));
+        return nullptr;
+    }
+    want.SetParam(UISERVICEHOSTPROXY_KEY, stub->AsObject());
+    return connection;
+}
+
 napi_value JsUIExtensionContext::OnConnectUIServiceExtension(napi_env env, NapiCallbackInfo& info)
 {
     TAG_LOGI(AAFwkTag::UISERVC_EXT, "called");
@@ -1153,9 +1173,10 @@ napi_value JsUIExtensionContext::OnConnectUIServiceExtension(napi_env env, NapiC
         return result;
     }
 
-    sptr<JSUIServiceUIExtConnection> connection = sptr<JSUIServiceUIExtConnection>::MakeSptr(env);
-    sptr<UIExtensionServiceHostStubImpl> stub = connection->GetServiceHostStub();
-    want.SetParam(UISERVICEHOSTPROXY_KEY, stub->AsObject());
+    sptr<JSUIServiceUIExtConnection> connection = CreateUIServiceExtConnection(env, want);
+    if (connection == nullptr) {
+        return CreateJsUndefined(env);
+    }
 
     result = nullptr;
     std::unique_ptr<NapiAsyncTask> uasyncTask = CreateAsyncTaskWithLastParam(env, nullptr, nullptr, nullptr, &result);
@@ -1203,6 +1224,21 @@ void JsUIExtensionContext::DoConnectUIServiceExtension(napi_env env,
     }
 }
 
+static void DoDisconnectUIServiceExtensionComplete(napi_env env, NapiAsyncTask& task,
+    int64_t connectId, std::shared_ptr<ErrCode> innerErrCode)
+{
+    if (*innerErrCode == static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INVALID_CONTEXT)) {
+        task.Reject(env, CreateJsError(env, AbilityErrorCode::ERROR_CODE_INVALID_CONTEXT));
+        UIServiceConnection::RemoveUIServiceExtensionConnection(connectId);
+    } else if (*innerErrCode == static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER)) {
+        task.Reject(env, CreateJsError(env, static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER),
+            GetInnerErrorMsg(AbilityInnerErrorMsg::CONNECTION_NOT_FOUND)));
+        UIServiceConnection::RemoveUIServiceExtensionConnection(connectId);
+    } else {
+        task.ResolveWithNoError(env, CreateJsUndefined(env));
+    }
+}
+
 napi_value JsUIExtensionContext::OnDisconnectUIServiceExtension(napi_env env, NapiCallbackInfo& info)
 {
     if (info.argc < ARGC_ONE) {
@@ -1235,21 +1271,13 @@ napi_value JsUIExtensionContext::OnDisconnectUIServiceExtension(napi_env env, Na
         if (!connection) {
             TAG_LOGW(AAFwkTag::UISERVC_EXT, "null connection");
             *innerErrCode = static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER);
+            return;
         }
         context->DisconnectAbility(want, connection);
     };
     NapiAsyncTask::CompleteCallback complete =
         [connectId, innerErrCode](napi_env env, NapiAsyncTask& task, int32_t status) {
-            if (*innerErrCode == static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INVALID_CONTEXT)) {
-                task.Reject(env, CreateJsError(env, AbilityErrorCode::ERROR_CODE_INVALID_CONTEXT));
-                UIServiceConnection::RemoveUIServiceExtensionConnection(connectId);
-            } else if (*innerErrCode == static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER)) {
-                task.Reject(env, CreateJsError(env, static_cast<int32_t>(AbilityErrorCode::ERROR_CODE_INNER),
-                    GetInnerErrorMsg(AbilityInnerErrorMsg::CONNECTION_NOT_FOUND)));
-                UIServiceConnection::RemoveUIServiceExtensionConnection(connectId);
-            } else {
-                task.ResolveWithNoError(env, CreateJsUndefined(env));
-            }
+            DoDisconnectUIServiceExtensionComplete(env, task, connectId, innerErrCode);
         };
     napi_value result = nullptr;
     NapiAsyncTask::Schedule("JsUIExtensionContext::OnDisconnectUIServiceExtension",
@@ -1259,6 +1287,11 @@ napi_value JsUIExtensionContext::OnDisconnectUIServiceExtension(napi_env env, Na
 
 napi_value JsUIExtensionContext::OnReportDrawnCompleted(napi_env env, NapiCallbackInfo& info)
 {
+    if (info.argc < ARGC_ONE) {
+        TAG_LOGE(AAFwkTag::UI_EXT, "invalid argc");
+        ThrowTooFewParametersError(env);
+        return CreateJsUndefined(env);
+    }
     TAG_LOGD(AAFwkTag::UI_EXT, "called");
     auto innerErrorCode = std::make_shared<int32_t>(ERR_OK);
     NapiAsyncTask::ExecuteCallback execute = [weak = context_, innerErrorCode]() {

frameworks/native/ability/native/ui_service_extension_ability/js_ui_service_extension_context.cpp

@@ -667,6 +667,9 @@ void JSUIServiceExtensionConnection::HandleOnAbilityDisconnectDone(const AppExec
             });
         if (item != g_connects.end()) {
             // match bundlename && abilityname
+            if (item->second) {
+                item->second->RemoveConnectionObject();
+            }
             g_connects.erase(item);
             TAG_LOGD(
                 AAFwkTag::UISERVC_EXT, "OnAbilityDisconnectDone erase g_connects.size:%{public}zu", g_connects.size());

services/abilitymgr/src/ui_extension/ui_extension_ability_manager.cpp

@@ -580,15 +580,28 @@ int32_t UIExtensionAbilityManager::RegisterPreloadUIExtensionHostClient(const sp
 
     {
         std::lock_guard lock(preloadUIExtRecipientMapMutex_);
+        auto it = preloadUIExtensionHostClientDeathRecipients_.find(callerPid);
+        if (it != preloadUIExtensionHostClientDeathRecipients_.end()) {
+            TAG_LOGW(AAFwkTag::UI_EXT, "recipient added before, callerPid: %{public}d", callerPid);
+            return ERR_OK;
+        }
+        if (!callerToken->AddDeathRecipient(deathRecipient)) {
+            TAG_LOGE(AAFwkTag::UI_EXT, "AddDeathRecipient fail");
+            return INNER_ERR;
+        }
         preloadUIExtensionHostClientDeathRecipients_[callerPid] = deathRecipient;
     }
 
-    callerToken->AddDeathRecipient(deathRecipient);
     try {
         uiExtensionAbilityRecordMgr_->RegisterPreloadUIExtensionHostClient(callerToken);
     } catch (std::exception &e) {
         TAG_LOGE(AAFwkTag::UI_EXT, "RegisterPreloadUIExtensionHostClient failed, exception = %{public}s", e.what());
         callerToken->RemoveDeathRecipient(deathRecipient);
+        {
+            std::lock_guard lock(preloadUIExtRecipientMapMutex_);
+            preloadUIExtensionHostClientDeathRecipients_.erase(callerPid);
+        }
+        return INNER_ERR;
     }
     return ERR_OK;
 }
@@ -1550,7 +1563,6 @@ void UIExtensionAbilityManager::CompleteBackground(const std::shared_ptr<BaseExt
         return;
     }
     abilityRecord->SetAbilityState(AbilityState::BACKGROUND);
-    CHECK_POINTER(abilityRecord);
     auto sessionInfo = abilityRecord->GetSessionInfo();
     CHECK_POINTER(sessionInfo);
     TAG_LOGI(AAFwkTag::UI_EXT,

services/abilitymgr/src/want_receiver_stub.cpp

@@ -52,7 +52,7 @@ int WantReceiverStub::SendInner(MessageParcel &data, MessageParcel &reply)
 
 int WantReceiverStub::PerformReceiveInner(MessageParcel &data, MessageParcel &reply)
 {
-    Want *want = data.ReadParcelable<Want>();
+    std::unique_ptr<Want> want(data.ReadParcelable<Want>());
     if (want == nullptr) {
         TAG_LOGE(AAFwkTag::WANTAGENT, "null want");
         return ERR_INVALID_VALUE;
@@ -61,10 +61,9 @@ int WantReceiverStub::PerformReceiveInner(MessageParcel &data, MessageParcel &re
     int resultCode = data.ReadInt32();
     std::string bundleName = Str16ToStr8(data.ReadString16());
 
-    WantParams *wantParams = data.ReadParcelable<WantParams>();
+    std::unique_ptr<WantParams> wantParams(data.ReadParcelable<WantParams>());
     if (wantParams == nullptr) {
         TAG_LOGE(AAFwkTag::WANTAGENT, "null wantParams");
-        delete want;
         return ERR_INVALID_VALUE;
     }
 
@@ -72,8 +71,6 @@ int WantReceiverStub::PerformReceiveInner(MessageParcel &data, MessageParcel &re
     bool sticky = data.ReadBool();
     int sendingUser = data.ReadInt32();
     PerformReceive(*want, resultCode, bundleName, *wantParams, serialized, sticky, sendingUser);
-    delete want;
-    delete wantParams;
     return NO_ERROR;
 }
 }  // namespace AAFwk

services/abilitymgr/src/wants_info.cpp

@@ -15,6 +15,8 @@
 
 #include "wants_info.h"
 
+#include "hilog_tag_wrapper.h"
+
 namespace OHOS {
 namespace AAFwk {
 bool WantsInfo::ReadFromParcel(Parcel &parcel)
@@ -45,8 +47,14 @@ WantsInfo *WantsInfo::Unmarshalling(Parcel &parcel)
 
 bool WantsInfo::Marshalling(Parcel &parcel) const
 {
-    parcel.WriteParcelable(&want);
-    parcel.WriteString16(Str8ToStr16(resolvedTypes));
+    if (!parcel.WriteParcelable(&want)) {
+        TAG_LOGE(AAFwkTag::ABILITYMGR, "write want failed");
+        return false;
+    }
+    if (!parcel.WriteString16(Str8ToStr16(resolvedTypes))) {
+        TAG_LOGE(AAFwkTag::ABILITYMGR, "write resolvedTypes failed");
+        return false;
+    }
     return true;
 }
 }  // namespace AAFwk

test/unittest/BUILD.gn

@@ -555,6 +555,7 @@ group("unittest") {
       "ui_extension_utils_test:unittest",
       "ui_extension_ability_manager_test:unittest",
       "ui_extension_ability_manager_second_test:unittest",
+      "ui_extension_ability_manager_third_test:unittest",
       "update_caller_info_util_test:unittest",
       "uri_utils_second_test",
       "uri_utils_test:unittest",

test/unittest/js_ui_extension_context_test/BUILD.gn

@@ -28,13 +28,18 @@ ohos_unittest("js_ui_extension_context_test") {
     debug = false
   }
   include_dirs = [
+    "${ability_runtime_napi_path}/ability_manager",
     "${ability_runtime_path}/interfaces/kits/native/ability/native",
     "${ability_runtime_path}/interfaces/kits/native/ability/native/ability_runtime",
     "${ability_runtime_path}/interfaces/kits/native/ability/native/ui_extension_ability",
     "${ability_runtime_path}/interfaces/kits/native/ability/native/ui_extension_base",
+    "${ability_runtime_path}/interfaces/kits/native/ability/native/ui_service_extension_ability/connection",
   ]
 
-  sources = [ "js_ui_extension_context_test.cpp" ]
+  sources = [
+    "${ability_runtime_napi_path}/ability_manager/js_preload_ui_extension_callback_client.cpp",
+    "js_ui_extension_context_test.cpp",
+  ]
 
   configs = []
 
@@ -46,6 +51,7 @@ ohos_unittest("js_ui_extension_context_test") {
     "${ability_runtime_native_path}/ability/native:abilitykit_native",
     "${ability_runtime_native_path}/ability/native:extensionkit_native",
     "${ability_runtime_native_path}/ability/native:ui_extension",
+    "${ability_runtime_native_path}/ability/native:ui_service_extension_connection",
     "${ability_runtime_native_path}/ability:ability_context_native",
     "${ability_runtime_native_path}/appkit:app_context",
     "${ability_runtime_path}/js_environment/frameworks/js_environment:js_environment",

test/unittest/js_ui_extension_context_test/js_ui_extension_context_test.cpp

@@ -13,7 +13,9 @@
  * limitations under the License.
  */
 
+#include <atomic>
 #include <gtest/gtest.h>
+#include <limits>
 #include <singleton.h>
 #include <uv.h>
 #include "ability_context.h"
@@ -33,6 +35,10 @@
 #include "native_engine/native_engine.h"
 #include "js_runtime_lite.h"
 #include "napi_common_want.h"
+#include "js_preload_ui_extension_callback_client.h"
+#include "js_ui_service_proxy.h"
+#include "ui_extension_servicehost_stub_impl.h"
+#include "js_uiservice_uiext_connection.h"
 
 using namespace testing;
 using namespace testing::ext;
@@ -66,6 +72,7 @@ class MockDeferred : public NativeDeferred {
 public:
     void Resolve(napi_value data) override
     {
+        settled_ = true;
         resolved_ = true;
         if (nref_ != nullptr) {
             napi_delete_reference(env_, nref_);
@@ -78,10 +85,12 @@ public:
 
     void Reject(napi_value reason) override
     {
+        settled_ = true;
         resolved_ = false;
     }
 
 public:
+    static bool IsSettled() { return settled_; }
     static bool GetLastResolveStatus() { return resolved_; }
     static napi_ref GetLastResolveValue() { return nref_; }
     static void Clear()
@@ -90,10 +99,14 @@ public:
             delete (reinterpret_cast<NativeReference*>(nref_));
             nref_ = nullptr;
         }
+        settled_ = false;
+        resolved_ = false;
     }
+    static bool settled_;
     static bool resolved_;
     static napi_ref nref_;
 };
+bool MockDeferred::settled_ = false;
 bool MockDeferred::resolved_ = false;
 napi_ref MockDeferred::nref_ = nullptr;
 
@@ -122,8 +135,15 @@ public:
     virtual ErrCode DisconnectAbility(const AAFwk::Want &want,
         const sptr<AbilityConnectCallback> &connectCallback) const override
     {
+        disconnectAbilityCount_++;
         return ERR_OK;
     }
+
+    ErrCode ReportDrawnCompleted() override
+    {
+        reportDrawnCompletedCount_++;
+        return reportDrawnCompletedResult_;
+    }
 public:
     static void DoneConnect(int status)
     {
@@ -139,9 +159,14 @@ public:
         callback_->OnAbilityDisconnectDone(element, 0);
     }
     void SetConnectResult(ErrCode code) { connectRet_ = code; }
+    int32_t GetDisconnectAbilityCount() const { return disconnectAbilityCount_.load(); }
+    int32_t GetReportDrawnCompletedCount() const { return reportDrawnCompletedCount_.load(); }
 protected:
     static sptr<AbilityConnectCallback> callback_;
     ErrCode connectRet_ = ERR_OK;
+    ErrCode reportDrawnCompletedResult_ = ERR_OK;
+    mutable std::atomic<int32_t> disconnectAbilityCount_ = 0;
+    std::atomic<int32_t> reportDrawnCompletedCount_ = 0;
 };
 
 sptr<AbilityConnectCallback> MockAbilityContextImpl::callback_;
@@ -160,6 +185,7 @@ public:
     }
     void Connect(napi_value* argv, int32_t argc);
     void Disconnect(napi_value* argv, int32_t argc);
+    void ReportDrawnCompleted(napi_value* argv, int32_t argc);
 public:
     std::shared_ptr<JsUIExtensionContext> jsUIExtensionContext_;
     std::shared_ptr<MockAbilityContextImpl> abilityContextImpl_;
@@ -264,6 +290,27 @@ void UIExtensionContextTest::Disconnect(napi_value* argv, int32_t argc)
     }
 }
 
+void UIExtensionContextTest::ReportDrawnCompleted(napi_value* argv, int32_t argc)
+{
+    napi_callback func = [](napi_env env, napi_callback_info info) -> napi_value {
+        return JsUIExtensionContext::ReportDrawnCompleted(env, info);
+    };
+    HandleScope handleScope(env_);
+    napi_value recv = nullptr;
+    napi_create_object(env_, &recv);
+    napi_status wrapret = napi_wrap(env_, recv, jsUIExtensionContext_.get(),
+        [](napi_env env, void* data, void* hint) {}, nullptr, nullptr);
+    EXPECT_EQ(wrapret, napi_ok);
+
+    napi_value funcValue = nullptr;
+    napi_create_function(env_, "reportDrawnCompleted", NAPI_AUTO_LENGTH, func, nullptr, &funcValue);
+    napi_value funcResultValue = nullptr;
+    napi_status status = napi_call_function(env_, recv, funcValue, argc, argv, &funcResultValue);
+    if (status != napi_ok) {
+        TAG_LOGE(AAFwkTag::UI_EXT, "call reportDrawnCompleted failed %{public}d", status);
+    }
+}
+
 HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_0100, TestSize.Level1)
 {
     GTEST_LOG_(INFO) << "AbilityRuntime_UIExtensionContext_0100 start";
@@ -391,6 +438,95 @@ HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_0105, TestSiz
     GTEST_LOG_(INFO) << "AbilityRuntime_UIExtensionContext_0105 end";
 }
 
+/**
+ * @tc.name: AbilityRuntime_UIExtensionContext_DisconnectMissingConnection_0100
+ * @tc.desc: A missing connection is rejected without calling the native disconnect API.
+ * @tc.type: FUNC
+ */
+HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_DisconnectMissingConnection_0100, TestSize.Level1)
+{
+    HandleScope handleScope(env_);
+    TryCatch tryCatch(env_);
+    constexpr int64_t missingConnectionId = std::numeric_limits<int64_t>::max();
+    UIServiceConnection::RemoveUIServiceExtensionConnection(missingConnectionId);
+    sptr<IRemoteObject> remoteObject = nullptr;
+    napi_value proxy = AAFwk::JsUIServiceProxy::CreateJsUIServiceProxy(
+        env_, remoteObject, missingConnectionId, remoteObject);
+    ASSERT_NE(proxy, nullptr);
+    napi_value argv[] = { proxy };
+
+    Disconnect(argv, ARGC_ONE);
+    ArkNativeEngine* engine = reinterpret_cast<ArkNativeEngine*>(env_);
+    uv_loop_t* loop = engine->GetUVLoop();
+    RunNowait(loop);
+    RunNowait(loop);
+
+    EXPECT_FALSE(tryCatch.HasCaught());
+    EXPECT_EQ(abilityContextImpl_->GetDisconnectAbilityCount(), 0);
+    EXPECT_TRUE(MockDeferred::IsSettled());
+    EXPECT_FALSE(MockDeferred::GetLastResolveStatus());
+}
+
+/**
+ * @tc.name: AbilityRuntime_UIExtensionContext_ReportDrawnCompleted_0100
+ * @tc.desc: Calling reportDrawnCompleted without a callback reports too few parameters.
+ * @tc.type: FUNC
+ */
+HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_ReportDrawnCompleted_0100, TestSize.Level1)
+{
+    HandleScope handleScope(env_);
+    TryCatch tryCatch(env_);
+
+    ReportDrawnCompleted(nullptr, ARGC_ZERO);
+
+    EXPECT_TRUE(tryCatch.HasCaught());
+    EXPECT_EQ(abilityContextImpl_->GetReportDrawnCompletedCount(), 0);
+    tryCatch.ClearException();
+    ArkNativeEngine* engine = reinterpret_cast<ArkNativeEngine*>(env_);
+    if (!engine->lastException_.IsEmpty()) {
+        engine->lastException_.Empty();
+    }
+}
+
+/**
+ * @tc.name: AbilityRuntime_UIExtensionContext_ReportDrawnCompleted_0200
+ * @tc.desc: A valid callback invokes the native reportDrawnCompleted API and completes asynchronously.
+ * @tc.type: FUNC
+ */
+HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_ReportDrawnCompleted_0200, TestSize.Level1)
+{
+    HandleScope handleScope(env_);
+    bool callbackInvoked = false;
+    napi_value callback = nullptr;
+    ASSERT_EQ(napi_create_function(env_, "reportDrawnCallback", NAPI_AUTO_LENGTH,
+        MarkCallbackInvoked, &callbackInvoked, &callback), napi_ok);
+    napi_value argv[] = { callback };
+
+    ReportDrawnCompleted(argv, ARGC_ONE);
+    ArkNativeEngine* engine = reinterpret_cast<ArkNativeEngine*>(env_);
+    uv_loop_t* loop = engine->GetUVLoop();
+    RunNowait(loop);
+    RunNowait(loop);
+
+    EXPECT_EQ(abilityContextImpl_->GetReportDrawnCompletedCount(), 1);
+    EXPECT_TRUE(callbackInvoked);
+}
+
+/**
+ * @tc.name: AbilityRuntime_PreloadUIExtensionCallback_NullEnv_0100
+ * @tc.desc: Calling the preload callback client with a null NAPI environment does not crash.
+ * @tc.type: FUNC
+ */
+HWTEST_F(UIExtensionContextTest, AbilityRuntime_PreloadUIExtensionCallback_NullEnv_0100, TestSize.Level1)
+{
+    auto callbackClient = std::make_shared<JsPreloadUIExtensionCallbackClient>(nullptr, nullptr);
+    ASSERT_NE(callbackClient, nullptr);
+
+    callbackClient->CallJsPreloadedUIExtensionAbility(1);
+
+    SUCCEED();
+}
+
 HWTEST_F(UIExtensionContextTest, AbilityRuntime_UIExtensionContext_0106, TestSize.Level1)
 {
     GTEST_LOG_(INFO) << "AbilityRuntime_UIExtensionContext_0106 start";
@@ -1513,4 +1649,4 @@ HWTEST_F(UIExtensionContextTest, TerminateSelfWithResultEmbeddable_0400, TestSiz
     GTEST_LOG_(INFO) << "TerminateSelfWithResultEmbeddable_0400 end";
 }
 }  // namespace AAFwk
-}  // namespace OHOS
+}  // namespace OHOS

test/unittest/js_ui_service_extension_context_second_test/js_ui_service_extension_context_second_test.cpp

@@ -37,6 +37,17 @@ namespace OHOS {
 namespace AbilityRuntime {
 namespace {
 const int64_t  COMMECTION_ID = 100;
+
+napi_value MarkDisconnectCallbackInvoked(napi_env env, napi_callback_info info)
+{
+    size_t argc = 0;
+    void* data = nullptr;
+    napi_get_cb_info(env, info, &argc, nullptr, nullptr, &data);
+    if (data != nullptr) {
+        *static_cast<bool*>(data) = true;
+    }
+    return CreateJsUndefined(env);
+}
 }  // namespace
 
 class JsUiServiceExtensionContextSecondTest : public testing::Test {
@@ -71,6 +82,16 @@ void JsUiServiceExtensionContextSecondTest::SetUp()
 
 void JsUiServiceExtensionContextSecondTest::TearDown()
 {
+    {
+        std::lock_guard guard(g_connectsMutex);
+        for (auto &item : g_connects) {
+            if (item.second != nullptr) {
+                item.second->RemoveConnectionObject();
+            }
+        }
+        g_connects.clear();
+        g_serialNumber = 0;
+    }
     if (env_ != nullptr) {
         delete reinterpret_cast<NativeEngine*>(env_);
         env_ = nullptr;
@@ -167,6 +188,44 @@ HWTEST_F(JsUiServiceExtensionContextSecondTest, FindConnection_0100, TestSize.Le
     TAG_LOGI(AAFwkTag::TEST, "FindConnection_0100 end");
 }
 
+/**
+ * @tc.name: HandleOnAbilityDisconnectDone_0400
+ * @tc.desc: A matched connection releases its JS reference and is removed after disconnect.
+ * @tc.type: FUNC
+ */
+HWTEST_F(JsUiServiceExtensionContextSecondTest, HandleOnAbilityDisconnectDone_0400, TestSize.Level1)
+{
+    bool callbackInvoked = false;
+    napi_value connectionObject = nullptr;
+    ASSERT_EQ(napi_create_object(env_, &connectionObject), napi_ok);
+    napi_value onDisconnect = nullptr;
+    ASSERT_EQ(napi_create_function(env_, "onDisconnect", NAPI_AUTO_LENGTH,
+        MarkDisconnectCallbackInvoked, &callbackInvoked, &onDisconnect), napi_ok);
+    ASSERT_EQ(napi_set_named_property(env_, connectionObject, "onDisconnect", onDisconnect), napi_ok);
+
+    sptr<JSUIServiceExtensionConnection> connection = new JSUIServiceExtensionConnection(env_);
+    connection->SetJsConnectionObject(connectionObject);
+    connection->SetConnectionId(COMMECTION_ID);
+    AppExecFwk::ElementName element("device", "com.example.uiservice", "UIServiceExtensionAbility");
+    Want want;
+    want.SetElement(element);
+    ConnectionKey key;
+    key.want = want;
+    key.id = COMMECTION_ID;
+    key.accountId = -1;
+    {
+        std::lock_guard guard(g_connectsMutex);
+        g_connects.emplace(key, connection);
+    }
+
+    connection->HandleOnAbilityDisconnectDone(element, ERR_OK);
+
+    EXPECT_TRUE(callbackInvoked);
+    EXPECT_EQ(connection->jsConnectionObject_, nullptr);
+    std::lock_guard guard(g_connectsMutex);
+    EXPECT_TRUE(g_connects.empty());
+}
+
 /**
  * @tc.name: OnConnectServiceExtensionAbility_0100
  * @tc.desc: basic function test.
@@ -247,4 +306,4 @@ HWTEST_F(JsUiServiceExtensionContextSecondTest, OnDisConnectServiceExtensionAbil
     TAG_LOGI(AAFwkTag::TEST, "OnDisConnectServiceExtensionAbility_0200 end");
 }
 }  // namespace AbilityRuntime
-}  // namespace OHOS
+}  // namespace OHOS

test/unittest/ui_extension_ability_manager_test/ui_extension_ability_manager_test.cpp

@@ -1018,21 +1018,16 @@ HWTEST_F(UIExtensionAbilityManagerTest, AAFwk_AbilityMS_RegisterPreloadUIExtensi
 /*
  * Feature: UIExtensionAbilityManager
  * Function: RegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify RegisterPreloadUIExtensionHostClient with valid parameters
+ * CaseDescription: Verify registration fails when a local token cannot register a death recipient
  */
 HWTEST_F(UIExtensionAbilityManagerTest, AAFwk_AbilityMS_RegisterPreloadUIExtensionHostClient_003, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    std::shared_ptr<AbilityRecord> abilityRecord = serviceRecord_;
-    ASSERT_NE(abilityRecord, nullptr);
-    sptr<IRemoteObject> callerToken = abilityRecord->GetToken();
-    ASSERT_NE(callerToken, nullptr);
-    
-    int32_t res = connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
-    EXPECT_EQ(res, ERR_OK);
+    sptr<IRemoteObject> callerToken = serviceRecord_->GetToken();
+
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(callerToken), INNER_ERR);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
 }
 
 /*
@@ -1073,43 +1068,30 @@ HWTEST_F(UIExtensionAbilityManagerTest, AAFwk_AbilityMS_UnRegisterPreloadUIExten
 /*
  * Feature: UIExtensionAbilityManager
  * Function: UnRegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify UnRegisterPreloadUIExtensionHostClient with valid callerPid
+ * CaseDescription: Verify unregistering an unknown process is idempotent
  */
 HWTEST_F(UIExtensionAbilityManagerTest, AAFwk_AbilityMS_UnRegisterPreloadUIExtensionHostClient_003, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    int32_t callerPid = 5678;
-    
-    int32_t res = connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid);
-    EXPECT_EQ(res, ERR_OK);
+
+    EXPECT_EQ(connectManager->UnRegisterPreloadUIExtensionHostClient(5678), ERR_OK);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
 }
 
 /*
  * Feature: UIExtensionAbilityManager
  * Function: UnRegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify UnRegisterPreloadUIExtensionHostClient
+ * CaseDescription: Verify failed registration leaves no state for unregistering
  */
 HWTEST_F(UIExtensionAbilityManagerTest, AAFwk_AbilityMS_UnRegisterPreloadUIExtensionHostClient_004, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    std::shared_ptr<AbilityRecord> abilityRecord = serviceRecord_;
-    ASSERT_NE(abilityRecord, nullptr);
-    sptr<IRemoteObject> callerToken = abilityRecord->GetToken();
-    int32_t callerPid = IPCSkeleton::GetCallingPid();
-    connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 1);
-    
-    connectManager->UnRegisterPreloadUIExtensionHostClient(1);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 1);
+    sptr<IRemoteObject> callerToken = serviceRecord_->GetToken();
 
-    connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 0);
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(callerToken), INNER_ERR);
+    EXPECT_EQ(connectManager->UnRegisterPreloadUIExtensionHostClient(IPCSkeleton::GetCallingPid()), ERR_OK);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
 }
 
 /*
@@ -1340,21 +1322,16 @@ HWTEST_F(UIExtensionAbilityManagerTest, RegisterPreloadUIExtensionHostClient_002
 /*
  * Feature: UIExtensionAbilityManager
  * Function: RegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify RegisterPreloadUIExtensionHostClient with valid parameters
+ * CaseDescription: Verify registration fails when a local token cannot register a death recipient
  */
 HWTEST_F(UIExtensionAbilityManagerTest, RegisterPreloadUIExtensionHostClient_003, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    std::shared_ptr<AbilityRecord> abilityRecord = serviceRecord_;
-    ASSERT_NE(abilityRecord, nullptr);
-    sptr<IRemoteObject> callerToken = abilityRecord->GetToken();
-    ASSERT_NE(callerToken, nullptr);
-    
-    int32_t res = connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
-    EXPECT_EQ(res, ERR_OK);
+    sptr<IRemoteObject> callerToken = serviceRecord_->GetToken();
+
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(callerToken), INNER_ERR);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
 }
 
 /*
@@ -1395,43 +1372,30 @@ HWTEST_F(UIExtensionAbilityManagerTest, UnRegisterPreloadUIExtensionHostClient_0
 /*
  * Feature: UIExtensionAbilityManager
  * Function: UnRegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify UnRegisterPreloadUIExtensionHostClient with valid callerPid
+ * CaseDescription: Verify unregistering an unknown process is idempotent
  */
 HWTEST_F(UIExtensionAbilityManagerTest, UnRegisterPreloadUIExtensionHostClient_003, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    int32_t callerPid = 5678;
-    
-    int32_t res = connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid);
-    EXPECT_EQ(res, ERR_OK);
+
+    EXPECT_EQ(connectManager->UnRegisterPreloadUIExtensionHostClient(5678), ERR_OK);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
 }
 
 /*
  * Feature: UIExtensionAbilityManager
  * Function: UnRegisterPreloadUIExtensionHostClient
- * SubFunction: NA
- * FunctionPoints: NA
- * EnvConditions: NA
- * CaseDescription: Verify UnRegisterPreloadUIExtensionHostClient
+ * CaseDescription: Verify failed registration leaves no state for unregistering
  */
 HWTEST_F(UIExtensionAbilityManagerTest, UnRegisterPreloadUIExtensionHostClient_004, TestSize.Level1)
 {
     std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
-    std::shared_ptr<AbilityRecord> abilityRecord = serviceRecord_;
-    ASSERT_NE(abilityRecord, nullptr);
-    sptr<IRemoteObject> callerToken = abilityRecord->GetToken();
-    int32_t callerPid = IPCSkeleton::GetCallingPid();
-    connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 1);
-    
-    connectManager->UnRegisterPreloadUIExtensionHostClient(1);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 1);
+    sptr<IRemoteObject> callerToken = serviceRecord_->GetToken();
 
-    connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid);
-    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.size(), 0);
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(callerToken), INNER_ERR);
+    EXPECT_EQ(connectManager->UnRegisterPreloadUIExtensionHostClient(IPCSkeleton::GetCallingPid()), ERR_OK);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
 }
 
 /*

test/unittest/ui_extension_ability_manager_third_test/BUILD.gn

@@ -0,0 +1,101 @@
+# Copyright (c) 2026 Huawei Device Co., Ltd.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import("//build/test.gni")
+import("//foundation/ability/ability_runtime/ability_runtime.gni")
+
+module_output_path = "ability_runtime/ability_runtime/ui_extension_ability_manager_third"
+
+ohos_unittest("ui_extension_ability_manager_third_test") {
+  module_out_path = module_output_path
+  sanitize = {
+    cfi = true
+    cfi_cross_dso = true
+    debug = false
+    blocklist = "../../cfi_blocklist.txt"
+  }
+  branch_protector_ret = "pac_ret"
+
+  include_dirs = [
+    "${ability_runtime_test_path}/mock/services_abilitymgr_test/libs/system_ability_mock",
+    "${ability_runtime_test_path}/mock/frameworks_kits_ability_native_test/include",
+    "${ability_runtime_test_path}/mock/mock_sa_call",
+    "${ability_runtime_test_path}/mock/task_handler_wrap_mock/include",
+  ]
+
+  sources = [
+    # add mock file
+    "${ability_runtime_test_path}/mock/services_abilitymgr_test/libs/appexecfwk_core/src/appmgr/mock_app_scheduler.cpp",
+    "${ability_runtime_test_path}/mock/task_handler_wrap_mock/src/mock_task_handler_wrap.cpp",
+    "ui_extension_ability_manager_third_test.cpp",
+  ]
+
+  configs = [
+    "${ability_runtime_services_path}/abilitymgr:abilityms_config",
+    "${ability_runtime_test_path}/mock/services_abilitymgr_test:aafwk_mock_config",
+  ]
+  cflags = []
+  if (target_cpu == "arm") {
+    cflags += [ "-DBINDER_IPC_32BIT" ]
+  }
+  deps = [
+    "${ability_runtime_innerkits_path}/ability_manager:ability_connect_callback_stub",
+    "${ability_runtime_innerkits_path}/ability_manager:ability_manager",
+    "${ability_runtime_innerkits_path}/deps_wrapper:ability_deps_wrapper",
+    "${ability_runtime_native_path}/ability/native:abilitykit_native",
+    "${ability_runtime_services_path}/abilitymgr:abilityms",
+    "${ability_runtime_services_path}/common:perm_verification",
+    "${ability_runtime_services_path}/common:task_handler_wrap",
+    "${ability_runtime_test_path}/mock/services_abilitymgr_test/libs/aakit:aakit_mock",
+    "${ability_runtime_test_path}/mock/services_abilitymgr_test/libs/appexecfwk_core:appexecfwk_bundlemgr_mock",
+  ]
+
+  external_deps = [
+    "ability_base:want",
+    "ability_base:zuri",
+    "access_token:libaccesstoken_sdk",
+    "access_token:libnativetoken",
+    "access_token:libtoken_setproc",
+    "c_utils:utils",
+    "common_event_service:cesfwk_innerkits",
+    "eventhandler:libeventhandler",
+    "ffrt:libffrt",
+    "googletest:gmock_main",
+    "googletest:gtest_main",
+    "hilog:libhilog",
+    "hisysevent:libhisysevent",
+    "init:libbeget_proxy",
+    "ipc:ipc_core",
+    "napi:ace_napi",
+    "safwk:system_ability_fwk",
+    "samgr:samgr_proxy",
+    "selinux_adapter:librestorecon",
+  ]
+
+  if (ability_runtime_graphics) {
+    external_deps += [
+      "image_framework:image_native",
+      "window_manager:libwsutils",
+      "window_manager:scene_session",
+    ]
+  }
+
+  if (background_task_mgr_continuous_task_enable) {
+    external_deps += [ "background_task_mgr:bgtaskmgr_innerkits" ]
+  }
+}
+
+group("unittest") {
+  testonly = true
+  deps = [ ":ui_extension_ability_manager_third_test" ]
+}

test/unittest/ui_extension_ability_manager_third_test/ui_extension_ability_manager_third_test.cpp

@@ -0,0 +1,189 @@
+/*
+ * Copyright (c) 2026 Huawei Device Co., Ltd.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <memory>
+#include <string>
+#include <vector>
+
+#define private public
+#define protected public
+#include "ui_extension_ability_manager.h"
+#undef private
+#undef protected
+
+#include "ability_manager_errors.h"
+#include "errors.h"
+#include "ipc_skeleton.h"
+#include "iremote_object.h"
+#include "message_option.h"
+#include "message_parcel.h"
+
+using namespace testing::ext;
+
+namespace OHOS {
+namespace AAFwk {
+class MockPreloadHostClient final : public IRemoteObject {
+public:
+    explicit MockPreloadHostClient(bool addDeathRecipientResult)
+        : IRemoteObject(u"mock_preload_host_client"), addDeathRecipientResult_(addDeathRecipientResult)
+    {}
+
+    ~MockPreloadHostClient() override = default;
+
+    int32_t GetObjectRefCount() override
+    {
+        return 0;
+    }
+
+    int SendRequest(uint32_t code, MessageParcel &data, MessageParcel &reply, MessageOption &option) override
+    {
+        (void)code;
+        (void)data;
+        (void)reply;
+        (void)option;
+        return ERR_OK;
+    }
+
+    bool IsProxyObject() const override
+    {
+        return true;
+    }
+
+    bool CheckObjectLegality() const override
+    {
+        return true;
+    }
+
+    bool AddDeathRecipient(const sptr<DeathRecipient> &recipient) override
+    {
+        addDeathRecipientCount_++;
+        deathRecipient_ = recipient;
+        return addDeathRecipientResult_;
+    }
+
+    bool RemoveDeathRecipient(const sptr<DeathRecipient> &recipient) override
+    {
+        removeDeathRecipientCount_++;
+        return recipient == deathRecipient_;
+    }
+
+    bool Marshalling(Parcel &parcel) const override
+    {
+        (void)parcel;
+        return true;
+    }
+
+    sptr<IRemoteBroker> AsInterface() override
+    {
+        return nullptr;
+    }
+
+    int Dump(int fd, const std::vector<std::u16string> &args) override
+    {
+        (void)fd;
+        (void)args;
+        return ERR_OK;
+    }
+
+    bool addDeathRecipientResult_ = true;
+    int32_t addDeathRecipientCount_ = 0;
+    int32_t removeDeathRecipientCount_ = 0;
+    sptr<DeathRecipient> deathRecipient_ = nullptr;
+};
+
+class UIExtensionAbilityManagerThirdTest : public testing::Test {};
+
+/*
+ * Feature: UIExtensionAbilityManager
+ * Function: RegisterPreloadUIExtensionHostClient
+ * CaseDescription: Verify successful registration and unregister cleanup
+ */
+HWTEST_F(UIExtensionAbilityManagerThirdTest, RegisterPreloadUIExtensionHostClient_006, TestSize.Level1)
+{
+    std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
+    sptr<MockPreloadHostClient> callerToken = new MockPreloadHostClient(true);
+    const int32_t callerPid = IPCSkeleton::GetCallingPid();
+
+    int32_t res = connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
+
+    EXPECT_EQ(res, ERR_OK);
+    EXPECT_EQ(callerToken->addDeathRecipientCount_, 1);
+    EXPECT_EQ(connectManager->preloadUIExtensionHostClientDeathRecipients_.count(callerPid), 1);
+    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.count(
+        callerPid), 1);
+
+    res = connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid + 1);
+    EXPECT_EQ(res, ERR_OK);
+    EXPECT_EQ(callerToken->removeDeathRecipientCount_, 0);
+    EXPECT_EQ(connectManager->preloadUIExtensionHostClientDeathRecipients_.count(callerPid), 1);
+    EXPECT_EQ(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.count(
+        callerPid), 1);
+
+    res = connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid);
+    EXPECT_EQ(res, ERR_OK);
+    EXPECT_EQ(callerToken->removeDeathRecipientCount_, 1);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
+}
+
+/*
+ * Feature: UIExtensionAbilityManager
+ * Function: RegisterPreloadUIExtensionHostClient
+ * CaseDescription: Verify registration rollback when adding a death recipient fails
+ */
+HWTEST_F(UIExtensionAbilityManagerThirdTest, RegisterPreloadUIExtensionHostClient_007, TestSize.Level1)
+{
+    std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
+    sptr<MockPreloadHostClient> callerToken = new MockPreloadHostClient(false);
+
+    int32_t res = connectManager->RegisterPreloadUIExtensionHostClient(callerToken);
+
+    EXPECT_EQ(res, INNER_ERR);
+    EXPECT_EQ(callerToken->addDeathRecipientCount_, 1);
+    EXPECT_EQ(callerToken->removeDeathRecipientCount_, 0);
+    EXPECT_TRUE(connectManager->preloadUIExtensionHostClientDeathRecipients_.empty());
+    EXPECT_TRUE(connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.empty());
+}
+
+/*
+ * Feature: UIExtensionAbilityManager
+ * Function: RegisterPreloadUIExtensionHostClient
+ * CaseDescription: Verify duplicate registration for the same process is idempotent
+ */
+HWTEST_F(UIExtensionAbilityManagerThirdTest, RegisterPreloadUIExtensionHostClient_008, TestSize.Level1)
+{
+    std::shared_ptr<UIExtensionAbilityManager> connectManager = std::make_shared<UIExtensionAbilityManager>(0);
+    sptr<MockPreloadHostClient> firstCallerToken = new MockPreloadHostClient(true);
+    sptr<MockPreloadHostClient> secondCallerToken = new MockPreloadHostClient(true);
+    const int32_t callerPid = IPCSkeleton::GetCallingPid();
+
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(firstCallerToken), ERR_OK);
+    EXPECT_EQ(connectManager->RegisterPreloadUIExtensionHostClient(secondCallerToken), ERR_OK);
+
+    EXPECT_EQ(firstCallerToken->addDeathRecipientCount_, 1);
+    EXPECT_EQ(secondCallerToken->addDeathRecipientCount_, 0);
+    EXPECT_EQ(connectManager->preloadUIExtensionHostClientDeathRecipients_.size(), 1);
+    auto tokenIter = connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.find(
+        callerPid);
+    ASSERT_NE(tokenIter,
+        connectManager->uiExtensionAbilityRecordMgr_->preloadUIExtensionHostClientCallerTokens_.end());
+    sptr<IRemoteObject> expectedCallerToken = firstCallerToken;
+    EXPECT_EQ(tokenIter->second, expectedCallerToken);
+
+    EXPECT_EQ(connectManager->UnRegisterPreloadUIExtensionHostClient(callerPid), ERR_OK);
+}
+}  // namespace AAFwk
+}  // namespace OHOS

test/unittest/wants_info_test/BUILD.gn

@@ -33,6 +33,7 @@ ohos_unittest("wants_info_test") {
   external_deps = [
     "ability_base:want",
     "c_utils:utils",
+    "hilog:libhilog",
     "ipc:ipc_core",
   ]
 }

test/unittest/wants_info_test/wants_info_test.cpp

@@ -13,7 +13,9 @@
  * limitations under the License.
  */
 
+#include <cstdlib>
 #include <gtest/gtest.h>
+
 #include "parcel.h"
 #define private public
 #define protected public
@@ -28,8 +30,38 @@ using OHOS::AppExecFwk::ElementName;
 
 namespace OHOS {
 namespace AAFwk {
-#define SLEEP(milli) std::this_thread::sleep_for(std::chrono::seconds(milli))
-namespace {}  // namespace
+namespace {
+class LimitedAllocator final : public Allocator {
+public:
+    explicit LimitedAllocator(size_t maxAllocationSize) : maxAllocationSize_(maxAllocationSize) {}
+
+    ~LimitedAllocator() override = default;
+
+    void *Realloc(void *data, size_t newSize) override
+    {
+        if (newSize > maxAllocationSize_) {
+            return nullptr;
+        }
+        return std::realloc(data, newSize);
+    }
+
+    void *Alloc(size_t size) override
+    {
+        if (size > maxAllocationSize_) {
+            return nullptr;
+        }
+        return std::malloc(size);
+    }
+
+    void Dealloc(void *data) override
+    {
+        std::free(data);
+    }
+
+private:
+    size_t maxAllocationSize_;
+};
+}  // namespace
 class WantsInfoTest : public testing::Test {
 public:
     static void SetUpTestCase();
@@ -66,12 +98,49 @@ HWTEST_F(WantsInfoTest, WantsInfoTest_0100, TestSize.Level1)
     info.want = want;
     info.resolvedTypes = "nihao";
     Parcel parcel;
-    info.Marshalling(parcel);
+    ASSERT_TRUE(info.Marshalling(parcel));
     auto unInfo = WantsInfo::Unmarshalling(parcel);
+    ASSERT_NE(unInfo, nullptr);
     EXPECT_EQ(unInfo->want.GetElement().GetBundleName(), "com.ix.hiMusic");
     EXPECT_EQ(unInfo->want.GetElement().GetAbilityName(), "MusicSAbility");
     EXPECT_EQ(unInfo->resolvedTypes, "nihao");
     delete unInfo;
 }
+
+/*
+ * @tc.number    : WantsInfoTest_0200
+ * @tc.name      : Marshalling want failure
+ * @tc.desc      : Marshalling returns false when the Want cannot be written.
+ */
+HWTEST_F(WantsInfoTest, WantsInfoTest_0200, TestSize.Level1)
+{
+    WantsInfo info;
+    Parcel parcel(new LimitedAllocator(0));
+
+    EXPECT_FALSE(info.Marshalling(parcel));
+}
+
+/*
+ * @tc.number    : WantsInfoTest_0300
+ * @tc.name      : Marshalling resolvedTypes failure
+ * @tc.desc      : Marshalling returns false when resolvedTypes cannot be written after the Want.
+ */
+HWTEST_F(WantsInfoTest, WantsInfoTest_0300, TestSize.Level1)
+{
+    WantsInfo info;
+    Want want;
+    ElementName element("device", "com.ix.hiMusic", "MusicSAbility");
+    want.SetElement(element);
+    info.want = want;
+
+    Parcel wantParcel;
+    ASSERT_TRUE(wantParcel.WriteParcelable(&info.want));
+    const size_t wantParcelCapacity = wantParcel.GetDataCapacity();
+    ASSERT_GT(wantParcelCapacity, 0);
+    info.resolvedTypes.assign(wantParcelCapacity, 'a');
+    Parcel parcel(new LimitedAllocator(wantParcelCapacity));
+
+    EXPECT_FALSE(info.Marshalling(parcel));
+}
 }  // namespace AAFwk
 }  // namespace OHOS