PIR: Logic for handling condition and resuming actions

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/PirRunStateHandler.kt

@@ -38,13 +38,7 @@ import com.duckduckgo.pir.impl.models.ExtractedProfile
 import com.duckduckgo.pir.impl.pixels.PirPixelSender
 import com.duckduckgo.pir.impl.scheduling.JobRecordUpdater
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ClickResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExpectationResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExtractedResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.FillFormResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.GetCaptchaInfoResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.NavigateResponse
-import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.SolveCaptchaResponse
 import com.duckduckgo.pir.impl.store.PirEventsRepository
 import com.duckduckgo.pir.impl.store.PirRepository
 import com.duckduckgo.pir.impl.store.PirSchedulingRepository
@@ -56,10 +50,9 @@ import com.duckduckgo.pir.impl.store.db.EmailConfirmationEventType.EMAIL_CONFIRM
 import com.duckduckgo.pir.impl.store.db.PirBrokerScanLog
 import com.squareup.anvil.annotations.ContributesBinding
 import com.squareup.moshi.Moshi
-import com.squareup.moshi.adapters.PolymorphicJsonAdapterFactory
-import com.squareup.moshi.kotlin.reflect.KotlinJsonAdapterFactory
 import kotlinx.coroutines.withContext
 import javax.inject.Inject
+import javax.inject.Named
 
 interface PirRunStateHandler {
     suspend fun handleState(pirRunState: PirRunState)
@@ -170,24 +163,8 @@ class RealPirRunStateHandler @Inject constructor(
     private val jobRecordUpdater: JobRecordUpdater,
     private val pirSchedulingRepository: PirSchedulingRepository,
     private val currentTimeProvider: CurrentTimeProvider,
+    @Named("pir") private val moshi: Moshi,
 ) : PirRunStateHandler {
-    private val moshi: Moshi by lazy {
-        Moshi
-            .Builder()
-            .add(
-                PolymorphicJsonAdapterFactory
-                    .of(PirSuccessResponse::class.java, "actionType")
-                    .withSubtype(NavigateResponse::class.java, "navigate")
-                    .withSubtype(ExtractedResponse::class.java, "extract")
-                    .withSubtype(GetCaptchaInfoResponse::class.java, "getCaptchaInfo")
-                    .withSubtype(SolveCaptchaResponse::class.java, "solveCaptcha")
-                    .withSubtype(ClickResponse::class.java, "click")
-                    .withSubtype(ExpectationResponse::class.java, "expectation")
-                    .withSubtype(FillFormResponse::class.java, "fillForm"),
-            ).add(KotlinJsonAdapterFactory())
-            .build()
-    }
-
     private val pirSuccessAdapter by lazy { moshi.adapter(PirSuccessResponse::class.java) }
 
     override suspend fun handleState(pirRunState: PirRunState) =

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/actions/ConditionExpectationSucceededEventHandler.kt

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2025 DuckDuckGo
+ *
+ * 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.
+ */
+
+package com.duckduckgo.pir.impl.common.actions
+
+import com.duckduckgo.di.scopes.AppScope
+import com.duckduckgo.pir.impl.common.BrokerStepsParser.BrokerStep
+import com.duckduckgo.pir.impl.common.actions.EventHandler.Next
+import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event
+import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event.ConditionExpectationSucceeded
+import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event.ExecuteBrokerStepAction
+import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.State
+import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData.UserProfile
+import com.squareup.anvil.annotations.ContributesMultibinding
+import javax.inject.Inject
+import kotlin.reflect.KClass
+
+@ContributesMultibinding(
+    scope = AppScope::class,
+    boundType = EventHandler::class,
+)
+class ConditionExpectationSucceededEventHandler @Inject constructor() : EventHandler {
+    override val event: KClass<out Event> = ConditionExpectationSucceeded::class
+
+    override suspend fun invoke(
+        state: State,
+        event: Event,
+    ): Next {
+        val actionsToAppend = (event as ConditionExpectationSucceeded).conditionActions
+        val currentBrokerStep = state.brokerStepsToExecute[state.currentBrokerStepIndex]
+
+        val updatedBrokerSteps = state.brokerStepsToExecute.toMutableList()
+        val updatedBrokerActions = currentBrokerStep.actions.toMutableList().apply {
+            this.addAll(
+                state.currentActionIndex + 1,
+                actionsToAppend,
+            )
+        }
+        val updatedBrokerStep = when (currentBrokerStep) {
+            is BrokerStep.ScanStep -> currentBrokerStep.copy(actions = updatedBrokerActions)
+            is BrokerStep.OptOutStep -> currentBrokerStep.copy(actions = updatedBrokerActions)
+            is BrokerStep.EmailConfirmationStep -> currentBrokerStep.copy(actions = updatedBrokerActions)
+        }
+
+        updatedBrokerSteps[state.currentBrokerStepIndex] = updatedBrokerStep
+
+        return Next(
+            nextState = state.copy(
+                currentActionIndex = state.currentActionIndex + 1,
+                brokerStepsToExecute = updatedBrokerSteps,
+            ),
+            nextEvent = ExecuteBrokerStepAction(
+                UserProfile(
+                    userProfile = state.profileQuery,
+                ),
+            ),
+        )
+    }
+}

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/actions/JsActionSuccessEventHandler.kt

@@ -25,6 +25,7 @@ import com.duckduckgo.pir.impl.common.PirRunStateHandler.PirRunState.BrokerOptOu
 import com.duckduckgo.pir.impl.common.PirRunStateHandler.PirRunState.BrokerScanActionSucceeded
 import com.duckduckgo.pir.impl.common.actions.EventHandler.Next
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event
+import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event.ConditionExpectationSucceeded
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event.ExecuteBrokerStepAction
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.Event.JsActionSuccess
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.SideEffect.EvaluateJs
@@ -33,6 +34,7 @@ import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.SideEf
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.State
 import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData.UserProfile
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ClickResponse
+import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ConditionResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExpectationResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExtractedResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.FillFormResponse
@@ -151,6 +153,28 @@ class JsActionSuccessEventHandler @Inject constructor(
                     ),
                 )
             }
+
+            is ConditionResponse -> {
+                if (pirSuccessResponse.response.actions.isNotEmpty()) {
+                    Next(
+                        nextState = baseSuccessState,
+                        nextEvent = ConditionExpectationSucceeded(
+                            pirSuccessResponse.response.actions,
+                        ),
+                    )
+                } else {
+                    Next(
+                        nextState = baseSuccessState.copy(
+                            currentActionIndex = baseSuccessState.currentActionIndex + 1,
+                        ),
+                        nextEvent = ExecuteBrokerStepAction(
+                            UserProfile(
+                                userProfile = baseSuccessState.profileQuery,
+                            ),
+                        ),
+                    )
+                }
+            }
         }
     }
 }

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/actions/PirActionsRunnerStateEngine.kt

@@ -122,6 +122,10 @@ interface PirActionsRunnerStateEngine {
             val actionId: String,
             val responseData: ResponseData?,
         ) : Event()
+
+        data class ConditionExpectationSucceeded(
+            val conditionActions: List<BrokerAction>,
+        ) : Event()
     }
 
     /**

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/di/PirModule.kt

@@ -42,6 +42,7 @@ import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData.SolveCaptcha
 import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData.UserProfile
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ClickResponse
+import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ConditionResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExpectationResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.ExtractedResponse
 import com.duckduckgo.pir.impl.scripts.models.PirSuccessResponse.FillFormResponse
@@ -217,8 +218,7 @@ class PirModule {
                     .withSubtype(BrokerAction.SolveCaptcha::class.java, "solveCaptcha")
                     .withSubtype(BrokerAction.EmailConfirmation::class.java, "emailConfirmation")
                     .withSubtype(BrokerAction.Condition::class.java, "condition"),
-            )
-            .add(
+            ).add(
                 PolymorphicJsonAdapterFactory.of(BrokerStep::class.java, "stepType")
                     .withSubtype(ScanStep::class.java, "scan")
                     .withSubtype(OptOutStep::class.java, "optOut"),
@@ -230,7 +230,8 @@ class PirModule {
                     .withSubtype(SolveCaptchaResponse::class.java, "solveCaptcha")
                     .withSubtype(ClickResponse::class.java, "click")
                     .withSubtype(ExpectationResponse::class.java, "expectation")
-                    .withSubtype(FillFormResponse::class.java, "fillForm"),
+                    .withSubtype(FillFormResponse::class.java, "fillForm")
+                    .withSubtype(ConditionResponse::class.java, "condition"),
             )
             .add(KotlinJsonAdapterFactory())
             .build()

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/scripts/models/PirScriptResponseParams.kt

@@ -127,4 +127,14 @@ sealed class PirSuccessResponse(
     data class AdditionalData(
         val additionalData: String,
     )
+
+    data class ConditionResponse(
+        override val actionID: String,
+        override val actionType: String,
+        val response: ResponseData = ResponseData(),
+    ) : PirSuccessResponse(actionID, actionType) {
+        data class ResponseData(
+            val actions: List<BrokerAction> = emptyList(),
+        )
+    }
 }

pir/pir-impl/src/test/kotlin/com/duckduckgo/pir/impl/common/RealPirRunStateHandlerTest.kt

@@ -47,6 +47,7 @@ import com.duckduckgo.pir.impl.store.db.BrokerScanEventType.BROKER_SUCCESS
 import com.duckduckgo.pir.impl.store.db.EmailConfirmationEventType.EMAIL_CONFIRMATION_FAILED
 import com.duckduckgo.pir.impl.store.db.EmailConfirmationEventType.EMAIL_CONFIRMATION_SUCCESS
 import com.duckduckgo.pir.impl.store.db.PirBrokerScanLog
+import com.squareup.moshi.Moshi
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Rule
@@ -71,6 +72,7 @@ class RealPirRunStateHandlerTest {
     private val mockJobRecordUpdater: JobRecordUpdater = mock()
     private val mockSchedulingRepository: PirSchedulingRepository = mock()
     private val mockCurrentTimeProvider: CurrentTimeProvider = mock()
+    private val moshi: Moshi = Moshi.Builder().build()
 
     @Before
     fun setUp() {
@@ -83,6 +85,7 @@ class RealPirRunStateHandlerTest {
                 jobRecordUpdater = mockJobRecordUpdater,
                 pirSchedulingRepository = mockSchedulingRepository,
                 currentTimeProvider = mockCurrentTimeProvider,
+                moshi = moshi,
             )
 
         whenever(mockCurrentTimeProvider.currentTimeMillis()).thenReturn(testEventTimeInMillis)