feat(#1289): lightweight /fork via session.branch_whole (#1303)

This commit is contained in:
m57
2026-09-27 16:56:43 +03:00
committed by GitHub
parent 3aedb44bed
commit ef54804b72
5 changed files with 186 additions and 16 deletions
@@ -15,6 +15,7 @@ object WsMethods {
const val SESSION_DELETE = "session.delete" const val SESSION_DELETE = "session.delete"
const val SESSION_TITLE = "session.title" const val SESSION_TITLE = "session.title"
const val SESSION_BRANCH = "session.branch" const val SESSION_BRANCH = "session.branch"
const val SESSION_BRANCH_WHOLE = "session.branch_whole"
/** Replay recorded events newer than client's last-seen seq for a session. */ /** Replay recorded events newer than client's last-seen seq for a session. */
const val SESSION_EVENTS_SINCE = "session.events.since" const val SESSION_EVENTS_SINCE = "session.events.since"
@@ -483,7 +483,13 @@ class ChatViewModel(
val originalPending: PendingSend? = null, val originalPending: PendingSend? = null,
) )
private data class PendingBranchRequest(
val generation: Long,
val params: Map<String, Any>,
)
private val sessionRequestById = ConcurrentHashMap<String, SessionRequest>() private val sessionRequestById = ConcurrentHashMap<String, SessionRequest>()
private val branchWholeRequests = ConcurrentHashMap<String, PendingBranchRequest>()
private val outgoingRequestById = ConcurrentHashMap<String, String>() private val outgoingRequestById = ConcurrentHashMap<String, String>()
private val hardInterruptRequestById = ConcurrentHashMap<String, HardInterruptRequest>() private val hardInterruptRequestById = ConcurrentHashMap<String, HardInterruptRequest>()
private val queuedStagingIds = ConcurrentHashMap.newKeySet<String>() private val queuedStagingIds = ConcurrentHashMap.newKeySet<String>()
@@ -1663,6 +1669,7 @@ class ChatViewModel(
) { ) {
val method = idToMethod.remove(id) ?: return val method = idToMethod.remove(id) ?: return
val request = sessionRequestById.remove(id) val request = sessionRequestById.remove(id)
branchWholeRequests.remove(id)
if (request != null && isStaleSessionRequest(request)) { if (request != null && isStaleSessionRequest(request)) {
outgoingRequestById.remove(id)?.let { markPendingSend(it, PendingSendState.UNKNOWN) } outgoingRequestById.remove(id)?.let { markPendingSend(it, PendingSendState.UNKNOWN) }
hardInterruptRequestById.remove(id)?.let { markPendingSend(it.messageId, PendingSendState.UNKNOWN) } hardInterruptRequestById.remove(id)?.let { markPendingSend(it.messageId, PendingSendState.UNKNOWN) }
@@ -1749,7 +1756,9 @@ class ChatViewModel(
} }
} }
WsMethods.SESSION_BRANCH -> { WsMethods.SESSION_BRANCH,
WsMethods.SESSION_BRANCH_WHOLE,
-> {
val resultMap = result as? Map<String, Any?> ?: return val resultMap = result as? Map<String, Any?> ?: return
// The result carries BOTH ids: `session_id` is the runtime // The result carries BOTH ids: `session_id` is the runtime
// registry id, `stored_session_id` is the DB key. currentSessionId // registry id, `stored_session_id` is the DB key. currentSessionId
@@ -2086,6 +2095,7 @@ class ChatViewModel(
) { ) {
val method = idToMethod.remove(id) ?: return val method = idToMethod.remove(id) ?: return
val request = sessionRequestById.remove(id) val request = sessionRequestById.remove(id)
val pendingBranch = branchWholeRequests.remove(id)
if (request != null && isStaleSessionRequest(request)) { if (request != null && isStaleSessionRequest(request)) {
outgoingRequestById.remove(id)?.let { markPendingSend(it, PendingSendState.UNKNOWN) } outgoingRequestById.remove(id)?.let { markPendingSend(it, PendingSendState.UNKNOWN) }
hardInterruptRequestById.remove(id)?.let { markPendingSend(it.messageId, PendingSendState.UNKNOWN) } hardInterruptRequestById.remove(id)?.let { markPendingSend(it.messageId, PendingSendState.UNKNOWN) }
@@ -2165,6 +2175,26 @@ class ChatViewModel(
} }
} }
if (method == WsMethods.SESSION_BRANCH_WHOLE) {
val code =
(error as? JsonRpcError)?.code
?: (error as? Map<*, *>)?.get("code") as? Int
if (code == -32601 && pendingBranch != null && pendingBranch.generation == sessionGeneration) {
val params = pendingBranch.params
val generation = pendingBranch.generation
viewModelScope.launch(ioDispatcher) {
wsClient.send(
WsMethods.SESSION_BRANCH,
params,
onSent = { branchId ->
trackSessionRequest(branchId, WsMethods.SESSION_BRANCH, generation)
},
)
}
return
}
}
// Surface error in UI (these are server-pushed RpcError for // Surface error in UI (these are server-pushed RpcError for
// fire-and-forget RPCs — awaited RPCs handle their own failure // fire-and-forget RPCs — awaited RPCs handle their own failure
// via the HermesWsClient.request() deferred). // via the HermesWsClient.request() deferred).
@@ -3940,9 +3970,8 @@ class ChatViewModel(
} }
/** /**
* Fork the active conversation via the session.branch WS RPC (issue #533). * Fork the active conversation via the lightweight session.branch_whole WS RPC (issue #1289).
* The backend already supports session.branch; the mobile previously had * Falls back to session.branch if the gateway answers -32601 (unknown method).
* no client surface, so `/fork` fell through to command.dispatch and 4018'd.
* The optional arg becomes the new branch's title. * The optional arg becomes the new branch's title.
*/ */
private fun branchSession(command: String) { private fun branchSession(command: String) {
@@ -3957,9 +3986,12 @@ class ChatViewModel(
val generation = sessionGeneration val generation = sessionGeneration
viewModelScope.launch(ioDispatcher) { viewModelScope.launch(ioDispatcher) {
wsClient.send( wsClient.send(
WsMethods.SESSION_BRANCH, WsMethods.SESSION_BRANCH_WHOLE,
params, params,
onSent = { id -> trackSessionRequest(id, WsMethods.SESSION_BRANCH, generation) }, onSent = { id ->
branchWholeRequests[id] = PendingBranchRequest(generation, params)
trackSessionRequest(id, WsMethods.SESSION_BRANCH_WHOLE, generation)
},
) )
} }
} }
@@ -6006,6 +6038,7 @@ class ChatViewModel(
private fun forgetRequest(id: String) { private fun forgetRequest(id: String) {
idToMethod.remove(id) idToMethod.remove(id)
sessionRequestById.remove(id) sessionRequestById.remove(id)
branchWholeRequests.remove(id)
} }
// ── Search ──────────────────────────────────────────────────────────── // ── Search ────────────────────────────────────────────────────────────
@@ -10,7 +10,8 @@ package com.m57.hermescontrol.ui.chat
* `/fork` and `/model` are NOT sent via [SlashResult.RpcDispatch] (which maps * `/fork` and `/model` are NOT sent via [SlashResult.RpcDispatch] (which maps
* to the `command.dispatch` RPC — that RPC only knows quick/plugin/bundle/skill * to the `command.dispatch` RPC — that RPC only knows quick/plugin/bundle/skill
* commands and 4018s on everything else). They are real backend commands that * commands and 4018s on everything else). They are real backend commands that
* get their own results: `/fork` goes via the `session.branch` RPC and `/model` * get their own results: `/fork` goes via the `session.branch_whole` RPC (with
* fallback to `session.branch` on older gateways) and `/model`
* via the `config.set` RPC (key="model" → gateway `_apply_model_switch`; the * via the `config.set` RPC (key="model" → gateway `_apply_model_switch`; the
* TUI gateway's `prompt.submit` does NOT parse slash commands, so sending it as * TUI gateway's `prompt.submit` does NOT parse slash commands, so sending it as
* a normal prompt makes the LLM treat it as text). `/update` is intercepted * a normal prompt makes the LLM treat it as text). `/update` is intercepted
@@ -289,7 +289,7 @@ class WsProfileParamsTest {
* home) without being listed here, this test fails — the set cannot drift. * home) without being listed here, this test fails — the set cannot drift.
* *
* Methods that operate on an ALREADY-BOUND session (session.usage, * Methods that operate on an ALREADY-BOUND session (session.usage,
* session.interrupt, session.branch, session.redirect, * session.interrupt, session.branch, session.branch_whole, session.redirect,
* session.context_breakdown, prompt.submit, slash/command dispatch, * session.context_breakdown, prompt.submit, slash/command dispatch,
* approval/clarify/secret/sudo responses, file attach, process.*, * approval/clarify/secret/sudo responses, file attach, process.*,
* subscription.*, commands_catalog, usage.bars) resolve via the * subscription.*, commands_catalog, usage.bars) resolve via the
@@ -761,12 +761,12 @@ class ChatViewModelTest {
id id
} }
// /fork with an optional branch title. // /fork with an optional branch title sends session.branch_whole.
viewModel.sendMessage("/fork my-fork") viewModel.sendMessage("/fork my-fork")
advanceUntilIdle() advanceUntilIdle()
val branchSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH } val branchSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH_WHOLE }
assertNotNull("session.branch should be dispatched for /fork", branchSent) assertNotNull("session.branch_whole should be dispatched for /fork", branchSent)
assertEquals(sessionId, branchSent!!.second["session_id"]) assertEquals(sessionId, branchSent!!.second["session_id"])
assertEquals("my-fork", branchSent.second["name"]) assertEquals("my-fork", branchSent.second["name"])
} }
@@ -811,12 +811,147 @@ class ChatViewModelTest {
viewModel.sendMessage("/fork") viewModel.sendMessage("/fork")
advanceUntilIdle() advanceUntilIdle()
val branchSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH } val branchSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH_WHOLE }
assertNotNull("session.branch should be dispatched for /fork", branchSent) assertNotNull("session.branch_whole should be dispatched for /fork", branchSent)
assertEquals(sessionId, branchSent!!.second["session_id"]) assertEquals(sessionId, branchSent!!.second["session_id"])
assertFalse("name param should be omitted when no title given", branchSent.second.containsKey("name")) assertFalse("name param should be omitted when no title given", branchSent.second.containsKey("name"))
} }
@Test
fun testSlashCommand_fork_branchWholeSuccess_handledLikeSessionBranch() =
runTest {
val (viewModel, _) = createViewModelWithSession()
val captured = mutableListOf<Pair<String, Map<String, Any>>>()
every { HermesWsClient.send(any(), any(), any()) } answers {
val id = "req-${captured.size + 1}"
captured.add(arg<String>(0) to (arg<Map<String, Any>>(1)))
arg<((String) -> Unit)?>(2)?.invoke(id)
id
}
viewModel.sendMessage("/fork")
advanceUntilIdle()
val branchWholeSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH_WHOLE }
assertNotNull("session.branch_whole dispatched", branchWholeSent)
mockEventsFlow.emit(
WsEvent.RpcResult(
"req-1",
mapOf(
"session_id" to "runtime-branch-whole",
"stored_session_id" to "stored-branch-whole",
"title" to "Whole Branch",
),
),
)
advanceUntilIdle()
assertEquals("runtime-branch-whole", ActiveSessionHolder.activeSessionId.value)
assertEquals("stored-branch-whole", viewModel.uiState.value.currentSessionId)
assertTrue(
"System message 'Session branched' should be added",
viewModel.uiState.value.messages
.any { it.content.contains("Session branched") },
)
// No fallback was sent
assertEquals(0, captured.count { it.first == WsMethods.SESSION_BRANCH })
assertEquals(1, captured.count { it.first == WsMethods.SESSION_BRANCH_WHOLE })
}
@Test
fun testSlashCommand_fork_branchWholeError32601_retriesWithSessionBranch() =
runTest {
val (viewModel, sessionId) = createViewModelWithSession()
val captured = mutableListOf<Pair<String, Map<String, Any>>>()
every { HermesWsClient.send(any(), any(), any()) } answers {
val id = "req-${captured.size + 1}"
captured.add(arg<String>(0) to (arg<Map<String, Any>>(1)))
arg<((String) -> Unit)?>(2)?.invoke(id)
id
}
viewModel.sendMessage("/fork my-retry-branch")
advanceUntilIdle()
val branchWholeSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH_WHOLE }
assertNotNull("session.branch_whole should be sent", branchWholeSent)
assertEquals("my-retry-branch", branchWholeSent!!.second["name"])
assertEquals(sessionId, branchWholeSent.second["session_id"])
// Backend answers -32601 unknown method
mockEventsFlow.emit(
WsEvent.RpcError(
"req-1",
JsonRpcError(code = -32601, message = "unknown method: session.branch_whole"),
),
)
advanceUntilIdle()
// Retries ONCE with session.branch using same params
val branchSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH }
assertNotNull("session.branch fallback should be sent", branchSent)
assertEquals("my-retry-branch", branchSent!!.second["name"])
assertEquals(sessionId, branchSent.second["session_id"])
assertEquals(1, captured.count { it.first == WsMethods.SESSION_BRANCH })
// session.branch result is then handled
mockEventsFlow.emit(
WsEvent.RpcResult(
"req-2",
mapOf(
"session_id" to "runtime-branch-fallback",
"stored_session_id" to "stored-branch-fallback",
"title" to "my-retry-branch",
),
),
)
advanceUntilIdle()
assertEquals("runtime-branch-fallback", ActiveSessionHolder.activeSessionId.value)
assertEquals("stored-branch-fallback", viewModel.uiState.value.currentSessionId)
assertTrue(
"System message 'Session branched' should be added",
viewModel.uiState.value.messages
.any { it.content.contains("Session branched") },
)
}
@Test
fun testSlashCommand_fork_branchWholeOtherError_noFallback() =
runTest {
val (viewModel, _) = createViewModelWithSession()
val captured = mutableListOf<Pair<String, Map<String, Any>>>()
every { HermesWsClient.send(any(), any(), any()) } answers {
val id = "req-${captured.size + 1}"
captured.add(arg<String>(0) to (arg<Map<String, Any>>(1)))
arg<((String) -> Unit)?>(2)?.invoke(id)
id
}
viewModel.sendMessage("/fork")
advanceUntilIdle()
val branchWholeSent = captured.firstOrNull { it.first == WsMethods.SESSION_BRANCH_WHOLE }
assertNotNull("session.branch_whole should be sent", branchWholeSent)
// Backend answers with code other than -32601
mockEventsFlow.emit(
WsEvent.RpcError(
"req-1",
JsonRpcError(code = -32600, message = "invalid request"),
),
)
advanceUntilIdle()
// No retry/fallback sent
assertEquals(0, captured.count { it.first == WsMethods.SESSION_BRANCH })
assertTrue(
viewModel.uiState.value.errorMessage
?.contains("session.branch_whole") == true,
)
}
@Test @Test
fun testSlashCommand_stop_sendsInterrupt() = fun testSlashCommand_stop_sendsInterrupt() =
runTest { runTest {
@@ -5234,7 +5369,7 @@ class ChatViewModelTest {
viewModel.sendMessage("/fork") viewModel.sendMessage("/fork")
runCurrent() runCurrent()
val branchRequest = captured.last { it.first == WsMethods.SESSION_BRANCH }.second val branchRequest = captured.last { it.first == WsMethods.SESSION_BRANCH_WHOLE }.second
mockEventsFlow.emit( mockEventsFlow.emit(
WsEvent.RpcResult( WsEvent.RpcResult(
branchRequest, branchRequest,
@@ -6068,11 +6203,11 @@ class ChatViewModelTest {
.SessionMessagesResponse(messages = emptyList()), .SessionMessagesResponse(messages = emptyList()),
) )
// /fork sends session.branch keyed on the runtime id. // /fork sends session.branch_whole keyed on the runtime id.
viewModel.sendMessage("/fork") viewModel.sendMessage("/fork")
advanceUntilIdle() advanceUntilIdle()
val branchId = captured.last { it.first == WsMethods.SESSION_BRANCH }.second val branchId = captured.last { it.first == WsMethods.SESSION_BRANCH_WHOLE }.second
val branchStorage = "branch-storage-1" val branchStorage = "branch-storage-1"
mockEventsFlow.emit( mockEventsFlow.emit(
WsEvent.RpcResult( WsEvent.RpcResult(