Skip to content

Commit 76382a3

Browse files
YoussefKeyrouzYoussefKeyrouz
committed
Per-chat model and reasoning restoration on existing Duck.ai chats.
1 parent b943724 commit 76382a3

8 files changed

Lines changed: 766 additions & 56 deletions

File tree

‎duckchat/duckchat-impl/src/main/java/com/duckduckgo/duckchat/impl/models/DuckAiModelManager.kt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,15 @@ data class ModelState(
4545
val selectedModelShortName: String? = null,
4646
val userTier: UserTier = UserTier.FREE,
4747
val attachmentLimits: AttachmentLimits = AttachmentLimits(),
48+
/** User's persisted global reasoning mode. Used for new chats. */
4849
val selectedReasoningMode: ReasoningMode? = null,
4950
val availableReasoningModes: List<AvailableReasoningMode> = emptyList(),
51+
/** Reasoning mode for the current chat. Session only (not persisted)
52+
* Note: if we ever want to preserve picks across unsubmitted chats
53+
* (User selected a mode but navigated away or switched tab without submission)
54+
* change this to a bounded Map<String, ReasoningMode> keyed by chatId.
55+
* */
56+
val chatScopedReasoningMode: ReasoningMode? = null,
5057
)
5158

5259
interface DuckAiModelManager {
@@ -58,6 +65,8 @@ interface DuckAiModelManager {
5865

5966
suspend fun selectReasoningMode(mode: ReasoningMode)
6067

68+
fun setChatScopedReasoningMode(mode: ReasoningMode?)
69+
6170
fun getSelectedModelId(): String?
6271

6372
fun getResolvedReasoningEffort(): String?
@@ -207,6 +216,10 @@ class RealDuckAiModelManager @Inject constructor(
207216
}
208217
}
209218

219+
override fun setChatScopedReasoningMode(mode: ReasoningMode?) {
220+
_modelState.value = _modelState.value.copy(chatScopedReasoningMode = mode)
221+
}
222+
210223
override fun getSelectedModelId(): String? = _modelState.value.selectedModelId
211224

212225
override fun getResolvedReasoningEffort(): String? {

‎duckchat/duckchat-impl/src/main/java/com/duckduckgo/duckchat/impl/models/Reasoning.kt‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package com.duckduckgo.duckchat.impl.models
1818

19+
import com.duckduckgo.duckchat.store.impl.DuckAiChat
20+
1921
enum class ReasoningEffort(val rawValue: String) {
2022
NONE("none"),
2123
MINIMAL("minimal"),
@@ -101,4 +103,21 @@ object ReasoningResolver {
101103
mode: ReasoningMode?,
102104
available: List<AvailableReasoningMode>,
103105
): ReasoningEffort? = available.firstOrNull { it.mode == mode }?.effort
106+
107+
/** Reasoning data for [chat]; `null` if its model isn't in [modelState] (callers fall back to global). */
108+
internal fun forChat(chat: DuckAiChat, modelState: ModelState): ChatReasoningResolution? {
109+
val chatModel = modelState.models.firstOrNull { it.id == chat.model } ?: return null
110+
val available = availableModes(
111+
supported = chatModel.supportedReasoningEfforts,
112+
effortAccess = chatModel.reasoningEffortAccess,
113+
)
114+
val mode = modelState.chatScopedReasoningMode ?: ReasoningMode.from(chat.reasoningMode)
115+
return ChatReasoningResolution(available = available, mode = mode)
116+
}
104117
}
118+
119+
/** Resolved reasoning data for an existing chat: the rows that apply and the mode to use. */
120+
internal data class ChatReasoningResolution(
121+
val available: List<AvailableReasoningMode>,
122+
val mode: ReasoningMode?,
123+
)

‎duckchat/duckchat-impl/src/main/java/com/duckduckgo/duckchat/impl/ui/NativeInputModeWidgetViewModel.kt‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,12 @@ import com.duckduckgo.duckchat.impl.inputscreen.ui.InputScreenConfigResolver
5151
import com.duckduckgo.duckchat.impl.inputscreen.ui.suggestions.ChatSuggestion
5252
import com.duckduckgo.duckchat.impl.inputscreen.ui.suggestions.reader.ChatSuggestionsReader
5353
import com.duckduckgo.duckchat.impl.models.DuckAiModelManager
54+
import com.duckduckgo.duckchat.impl.models.ReasoningResolver
5455
import com.duckduckgo.duckchat.impl.nativeinput.NativeInputPlugin
5556
import com.duckduckgo.duckchat.impl.pixel.DuckChatPixelName
5657
import com.duckduckgo.duckchat.impl.store.DefaultTogglePosition
58+
import com.duckduckgo.duckchat.store.impl.DuckAiChat
59+
import com.duckduckgo.duckchat.store.impl.DuckAiChatStore
5760
import com.duckduckgo.subscriptions.api.Product
5861
import com.duckduckgo.subscriptions.api.Subscriptions
5962
import kotlinx.coroutines.CoroutineScope
@@ -67,6 +70,7 @@ import kotlinx.coroutines.flow.SharedFlow
6770
import kotlinx.coroutines.flow.SharingStarted
6871
import kotlinx.coroutines.flow.StateFlow
6972
import kotlinx.coroutines.flow.asStateFlow
73+
import kotlinx.coroutines.flow.collectLatest
7074
import kotlinx.coroutines.flow.combine
7175
import kotlinx.coroutines.flow.filterNotNull
7276
import kotlinx.coroutines.flow.firstOrNull
@@ -75,7 +79,6 @@ import kotlinx.coroutines.flow.receiveAsFlow
7579
import kotlinx.coroutines.flow.shareIn
7680
import kotlinx.coroutines.flow.update
7781
import kotlinx.coroutines.launch
78-
import logcat.logcat
7982
import javax.inject.Inject
8083

8184
data class ChatTabSuggestions(
@@ -100,6 +103,7 @@ class NativeInputModeWidgetViewModel @Inject constructor(
100103
private val nativeInputStatePublisher: NativeInputStatePublisher,
101104
private val nativeInputStateProvider: NativeInputStateProvider,
102105
private val modelManager: DuckAiModelManager,
106+
private val duckAiChatStore: DuckAiChatStore,
103107
@AppCoroutineScope private val appCoroutineScope: CoroutineScope,
104108
) : ViewModel() {
105109

@@ -124,13 +128,22 @@ class NativeInputModeWidgetViewModel @Inject constructor(
124128

125129
private val _chatId = MutableStateFlow<String?>(null)
126130

131+
private val currentChat = MutableStateFlow<DuckAiChat?>(null)
132+
127133
private val commandChannel = Channel<Command>(capacity = 1, onBufferOverflow = DROP_OLDEST)
128134
val commands = commandChannel.receiveAsFlow()
129135

130136
init {
131137
viewModelScope.launch {
132138
_plugins.value = nativeInputPlugins.getPlugins().toList()
133139
}
140+
viewModelScope.launch {
141+
_chatId.collectLatest { chatId ->
142+
modelManager.setChatScopedReasoningMode(null)
143+
currentChat.value = null
144+
if (chatId != null) currentChat.value = duckAiChatStore.getChatById(chatId)
145+
}
146+
}
134147
}
135148

136149
fun updatePluginContainerVisibility(isChatTab: Boolean) {
@@ -144,12 +157,26 @@ class NativeInputModeWidgetViewModel @Inject constructor(
144157
_modelPickerEnabled.value = enabled
145158
}
146159

160+
// currentChat can briefly hold the previous chat during a transition
161+
// Returns it only if it actually matches the current chatId.
162+
private fun validChat(): DuckAiChat? = currentChat.value?.takeIf { it.chatId == _chatId.value }
163+
147164
fun getSelectedModelId(): String? {
165+
// Existing chat: send chat's stored model.
166+
validChat()?.model?.let { return it }
167+
// New chat with picker off: nothing to send.
148168
if (!_modelPickerEnabled.value) return null
169+
// New chat with picker on: send what the user picked.
149170
return modelManager.getSelectedModelId()
150171
}
151172

152-
fun getResolvedReasoningEffort(): String? = modelManager.getResolvedReasoningEffort()
173+
fun getResolvedReasoningEffort(): String? {
174+
val chat = validChat() ?: return modelManager.getResolvedReasoningEffort()
175+
val resolution = ReasoningResolver.forChat(chat, modelManager.modelState.value)
176+
?: return modelManager.getResolvedReasoningEffort()
177+
val mode = ReasoningResolver.resolveMode(persisted = resolution.mode, available = resolution.available)
178+
return ReasoningResolver.effortFor(mode, resolution.available)?.rawValue
179+
}
153180

154181
fun getSelectedTool(): String? {
155182
val tabId = activeTabId.value ?: return null
@@ -211,15 +238,11 @@ class NativeInputModeWidgetViewModel @Inject constructor(
211238
}
212239
}
213240

214-
// Publish chatId to per-tab state. setActiveChatId can fire during widget attach before
215-
// configure() sets activeTabId — _chatId holds the value until then.
241+
// Publish chatId to per-tab state. Combine buffers it until activeTabId is set.
216242
viewModelScope.launch {
217243
combine(_chatId, activeTabId.filterNotNull()) { chatId, tabId -> tabId to chatId }
218244
.collect { (tabId, chatId) ->
219245
nativeInputStatePublisher.update(tabId) { it.copy(chatId = chatId) }
220-
// TODO: logs the chatId observers will see for this tab. Added for testing.
221-
// Remove once consumer logic is implemented.
222-
logcat { "Duck.ai native input: active chatId for tab=$tabId is now $chatId" }
223246
}
224247
}
225248
}

‎duckchat/duckchat-impl/src/main/java/com/duckduckgo/duckchat/impl/ui/nativeinput/views/ReasoningModePickerView.kt‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ import com.duckduckgo.duckchat.api.nativeinput.NativeInputState.InputContext
3939
import com.duckduckgo.duckchat.api.nativeinput.NativeInputStateProvider
4040
import com.duckduckgo.duckchat.impl.DuckChatConstants.DUCK_AI_FEATURE_PAGE
4141
import com.duckduckgo.duckchat.impl.R
42-
import com.duckduckgo.duckchat.impl.models.ModelState
4342
import com.duckduckgo.navigation.api.GlobalActivityStarter
4443
import com.duckduckgo.subscriptions.api.SubscriptionScreens.SubscriptionPurchase
4544
import com.duckduckgo.subscriptions.api.SubscriptionScreens.SubscriptionUpgrade
@@ -112,9 +111,11 @@ class ReasoningModePickerView @JvmOverloads constructor(
112111
stateJob?.cancel()
113112
stateJob = viewModel.state
114113
.onEach { state ->
115-
isVisible = state.availableReasoningModes.size > 1 &&
116-
state.availableReasoningModes.any { it.isAccessible }
117-
viewModel.resolvedMode(state)?.let { mode ->
114+
isVisible = state.visible
115+
// Popup is positioned by screen coordinates, not parented to the button. Dismiss
116+
// it explicitly when picker hides or stale rows stay tappable for the new chat.
117+
if (!state.visible) dismissPopup()
118+
state.displayedMode?.let { mode ->
118119
button.setImageResource(viewModel.iconResFor(mode))
119120
}
120121
}
@@ -143,8 +144,7 @@ class ReasoningModePickerView @JvmOverloads constructor(
143144

144145
private fun showMenu() {
145146
val state = viewModel.state.value
146-
if (state.availableReasoningModes.size <= 1) return
147-
if (state.availableReasoningModes.none { it.isAccessible }) return
147+
if (!state.visible) return
148148

149149
val container = LinearLayout(context).apply {
150150
orientation = LinearLayout.VERTICAL
@@ -168,8 +168,8 @@ class ReasoningModePickerView @JvmOverloads constructor(
168168
showAtPosition(popup)
169169
}
170170

171-
private fun populate(container: LinearLayout, popup: PopupWindow, state: ModelState) {
172-
viewModel.rows(state).forEach { row ->
171+
private fun populate(container: LinearLayout, popup: PopupWindow, state: ReasoningModePickerState) {
172+
state.rows.forEach { row ->
173173
val item = LayoutInflater.from(context)
174174
.inflate(R.layout.view_reasoning_mode_picker_item, container, false)
175175
item.findViewById<ImageView>(R.id.reasoningModeItemLeadingIcon)

‎duckchat/duckchat-impl/src/main/java/com/duckduckgo/duckchat/impl/ui/nativeinput/views/ReasoningModePickerViewModel.kt‎

Lines changed: 78 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,28 @@ import androidx.lifecycle.ViewModel
2222
import androidx.lifecycle.viewModelScope
2323
import com.duckduckgo.anvil.annotations.ContributesViewModel
2424
import com.duckduckgo.di.scopes.ViewScope
25+
import com.duckduckgo.duckchat.api.nativeinput.NativeInputState
26+
import com.duckduckgo.duckchat.api.nativeinput.NativeInputStateProvider
2527
import com.duckduckgo.duckchat.impl.R
2628
import com.duckduckgo.duckchat.impl.models.AvailableReasoningMode
2729
import com.duckduckgo.duckchat.impl.models.DuckAiModelManager
2830
import com.duckduckgo.duckchat.impl.models.ModelState
2931
import com.duckduckgo.duckchat.impl.models.ReasoningMode
3032
import com.duckduckgo.duckchat.impl.models.ReasoningResolver
33+
import com.duckduckgo.duckchat.store.impl.DuckAiChat
34+
import com.duckduckgo.duckchat.store.impl.DuckAiChatStore
3135
import kotlinx.coroutines.channels.BufferOverflow
3236
import kotlinx.coroutines.channels.Channel
3337
import kotlinx.coroutines.flow.Flow
38+
import kotlinx.coroutines.flow.MutableStateFlow
39+
import kotlinx.coroutines.flow.SharingStarted
3440
import kotlinx.coroutines.flow.StateFlow
41+
import kotlinx.coroutines.flow.collectLatest
42+
import kotlinx.coroutines.flow.combine
43+
import kotlinx.coroutines.flow.distinctUntilChanged
44+
import kotlinx.coroutines.flow.map
3545
import kotlinx.coroutines.flow.receiveAsFlow
46+
import kotlinx.coroutines.flow.stateIn
3647
import kotlinx.coroutines.launch
3748
import logcat.logcat
3849
import javax.inject.Inject
@@ -45,42 +56,96 @@ data class ReasoningModeRow(
4556
val selected: Boolean,
4657
)
4758

59+
/** Resolved snapshot the picker view renders from. */
60+
data class ReasoningModePickerState(
61+
val visible: Boolean,
62+
val rows: List<ReasoningModeRow>,
63+
val displayedMode: ReasoningMode?,
64+
)
65+
4866
@ContributesViewModel(ViewScope::class)
4967
class ReasoningModePickerViewModel @Inject constructor(
5068
private val modelManager: DuckAiModelManager,
69+
private val nativeInputStateProvider: NativeInputStateProvider,
70+
private val duckAiChatStore: DuckAiChatStore,
5171
) : ViewModel() {
5272

53-
val state: StateFlow<ModelState> = modelManager.modelState
73+
private val currentChat = MutableStateFlow<DuckAiChat?>(null)
74+
75+
init {
76+
viewModelScope.launch {
77+
// collectLatest: cancel in-flight lookup on chatId flip to avoid stale currentChat.
78+
nativeInputStateProvider.state
79+
.map { it.chatId }
80+
.distinctUntilChanged()
81+
.collectLatest { chatId ->
82+
currentChat.value = if (chatId == null) null else duckAiChatStore.getChatById(chatId)
83+
}
84+
}
85+
}
5486

55-
fun resolvedMode(state: ModelState): ReasoningMode? =
56-
ReasoningResolver.resolveMode(state.selectedReasoningMode, state.availableReasoningModes)
87+
/** Existing chat: rows derive from the chat's model. Otherwise: from the global selection. */
88+
val state: StateFlow<ReasoningModePickerState> = combine(
89+
modelManager.modelState,
90+
nativeInputStateProvider.state,
91+
currentChat,
92+
) { modelState, nativeState, chat ->
93+
resolveState(modelState, nativeState, chat)
94+
}.stateIn(
95+
scope = viewModelScope,
96+
started = SharingStarted.Eagerly,
97+
initialValue = ReasoningModePickerState(visible = false, rows = emptyList(), displayedMode = null),
98+
)
99+
100+
private fun resolveState(
101+
modelState: ModelState,
102+
nativeState: NativeInputState,
103+
chat: DuckAiChat?,
104+
): ReasoningModePickerState {
105+
val chatMatches = chat?.chatId == nativeState.chatId
106+
val chatResolution = chat?.let { ReasoningResolver.forChat(it, modelState) }
107+
if (nativeState.chatId != null && (!chatMatches || chatResolution == null)) {
108+
return ReasoningModePickerState(visible = false, rows = emptyList(), displayedMode = null)
109+
}
110+
val available = chatResolution?.available ?: modelState.availableReasoningModes
111+
val persistedMode = chatResolution?.mode ?: modelState.selectedReasoningMode
112+
val displayedMode = ReasoningResolver.resolveMode(persistedMode, available)
113+
val visible = available.size > 1 && available.any { it.isAccessible }
114+
val rows = available.map { it.toRow(selected = it.mode == displayedMode) }
115+
return ReasoningModePickerState(visible = visible, rows = rows, displayedMode = displayedMode)
116+
}
57117

58118
private val command = Channel<UpsellCommand>(capacity = 1, onBufferOverflow = BufferOverflow.DROP_OLDEST)
59119
val commands: Flow<UpsellCommand> = command.receiveAsFlow()
60120

61121
fun onModeTapped(mode: ReasoningMode, surface: PickerSurface) {
62-
val available = modelManager.modelState.value.availableReasoningModes.firstOrNull { it.mode == mode }
63-
if (available == null) {
122+
// Drop taps that race past the picker hiding (loading, mismatch, no accessible rows).
123+
if (!state.value.visible) return
124+
val chat = currentChat.value
125+
val chatResolution = chat?.let { ReasoningResolver.forChat(it, modelManager.modelState.value) }
126+
val available = chatResolution?.available ?: modelManager.modelState.value.availableReasoningModes
127+
128+
val match = available.firstOrNull { it.mode == mode }
129+
if (match == null) {
64130
logcat { "Duck.ai reasoning picker: tapped mode $mode not in available list, ignoring." }
65131
return
66132
}
67-
if (available.isAccessible) {
68-
viewModelScope.launch { modelManager.selectReasoningMode(mode) }
133+
if (match.isAccessible) {
134+
if (chatResolution != null) {
135+
modelManager.setChatScopedReasoningMode(mode)
136+
} else {
137+
viewModelScope.launch { modelManager.selectReasoningMode(mode) }
138+
}
69139
return
70140
}
71141
val userTier = modelManager.modelState.value.userTier
72-
val requiredTier = available.access?.requiredTier ?: run {
142+
val requiredTier = match.access?.requiredTier ?: run {
73143
logcat { "Duck.ai reasoning picker: gated mode $mode has no public required tier, ignoring." }
74144
return
75145
}
76146
routeUpsell(userTier, requiredTier, surface.origin)?.let { command.trySend(it) }
77147
}
78148

79-
fun rows(state: ModelState): List<ReasoningModeRow> {
80-
val resolved = resolvedMode(state)
81-
return state.availableReasoningModes.map { it.toRow(selected = it.mode == resolved) }
82-
}
83-
84149
@DrawableRes
85150
fun iconResFor(mode: ReasoningMode): Int = when (mode) {
86151
ReasoningMode.FAST -> R.drawable.ic_reasoning_fast_24

‎duckchat/duckchat-impl/src/test/kotlin/com/duckduckgo/duckchat/impl/models/RealDuckAiModelManagerTest.kt‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1131,6 +1131,28 @@ class RealDuckAiModelManagerTest {
11311131
verify(dataStore, never()).setSelectedReasoningMode(ReasoningMode.FAST.rawValue)
11321132
}
11331133

1134+
@Test
1135+
fun whenSetChatScopedReasoningModeThenStateUpdatedAndNotPersisted() = runTest {
1136+
testee = createManager()
1137+
1138+
testee.setChatScopedReasoningMode(ReasoningMode.REASONING)
1139+
1140+
assertEquals(ReasoningMode.REASONING, testee.modelState.value.chatScopedReasoningMode)
1141+
verify(dataStore, never()).setSelectedReasoningMode(any())
1142+
}
1143+
1144+
@Test
1145+
fun whenSetChatScopedReasoningModeToNullThenStateClearedAndNotPersisted() = runTest {
1146+
testee = createManager()
1147+
testee.setChatScopedReasoningMode(ReasoningMode.REASONING)
1148+
assertEquals(ReasoningMode.REASONING, testee.modelState.value.chatScopedReasoningMode)
1149+
1150+
testee.setChatScopedReasoningMode(null)
1151+
1152+
assertNull(testee.modelState.value.chatScopedReasoningMode)
1153+
verify(dataStore, never()).setSelectedReasoningMode(any())
1154+
}
1155+
11341156
@Test
11351157
fun whenSelectReasoningModeAndModeIsGatedThenIgnoredAndNotPersisted() = runTest {
11361158
// Defends against direct calls bypassing the picker's tap handler: gated modes must not be persistable.

0 commit comments

Comments
 (0)