Skip to content

Commit b7d0036

Browse files
OloloshechkinOloloshechkin
authored
Merge pull request #119 from nikita-ashihmin/nikita.ashihmin/fix-session-update-race
Do not report an existing session as missing in getOrCreateSessionHolder
2 parents bf9d47b + 80e3e30 commit b7d0036

2 files changed

Lines changed: 176 additions & 10 deletions

File tree

‎acp/src/commonMain/kotlin/com/agentclientprotocol/client/Client.kt‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -96,24 +96,29 @@ public class Client(
9696
* Creates a new entry only if there are some currently initializing sessions. Otherwise, throws in the case of missing session.
9797
*/
9898
private fun getOrCreateSessionHolder(sessionId: SessionId): ClientSessionHolder {
99-
val sessionsStorage = _sessions.value
100-
val holder = sessionsStorage.sessions[sessionId]
101-
if (holder != null) return holder
99+
// Fast path for the common case of an already registered session.
100+
_sessions.value.sessions[sessionId]?.let { return it }
102101
var clientSessionHolder: ClientSessionHolder? = null
102+
// Every branch below has to look the session up in `currentStorage` rather than rely on the read above:
103+
// `_sessions` can change between that read and the CAS, and it can change again between two attempts of
104+
// the CAS loop. A session that a concurrent newSession/loadSession registers in such a window - together
105+
// with decrementing `initializingSessionsCount` back to zero - would otherwise be reported as missing.
103106
_sessions.update { currentStorage ->
104-
if (currentStorage.initializingSessionsCount > 0) {
105-
val existingHolder = currentStorage.sessions[sessionId]
106-
if (existingHolder != null) {
107+
val existingHolder = currentStorage.sessions[sessionId]
108+
when {
109+
existingHolder != null -> {
107110
clientSessionHolder = existingHolder
108111
currentStorage
109-
} else {
112+
}
113+
currentStorage.initializingSessionsCount > 0 -> {
110114
val newHolder = ClientSessionHolder()
111115
clientSessionHolder = newHolder
112116
currentStorage.copy(sessions = currentStorage.sessions.put(sessionId, newHolder))
113117
}
114-
} else {
115-
clientSessionHolder = null
116-
currentStorage
118+
else -> {
119+
clientSessionHolder = null
120+
currentStorage
121+
}
117122
}
118123
}
119124
return clientSessionHolder ?: acpFail("Session $sessionId not found")

‎acp/src/jvmTest/kotlin/com/agentclientprotocol/client/ClientSessionUpdateDeliveryTest.kt‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
@file:OptIn(UnstableApi::class)
2+
3+
package com.agentclientprotocol.client
4+
5+
import com.agentclientprotocol.annotations.UnstableApi
6+
import com.agentclientprotocol.common.ClientSessionOperations
7+
import com.agentclientprotocol.common.SessionCreationParameters
8+
import com.agentclientprotocol.model.AcpMethod
9+
import com.agentclientprotocol.model.ContentBlock
10+
import com.agentclientprotocol.model.NewSessionResponse
11+
import com.agentclientprotocol.model.PermissionOption
12+
import com.agentclientprotocol.model.RequestPermissionResponse
13+
import com.agentclientprotocol.model.SessionId
14+
import com.agentclientprotocol.model.SessionNotification
15+
import com.agentclientprotocol.model.SessionUpdate
16+
import com.agentclientprotocol.protocol.Protocol
17+
import com.agentclientprotocol.rpc.ACPJson
18+
import com.agentclientprotocol.rpc.JsonRpcMessage
19+
import com.agentclientprotocol.rpc.JsonRpcNotification
20+
import com.agentclientprotocol.rpc.JsonRpcRequest
21+
import com.agentclientprotocol.rpc.JsonRpcResponse
22+
import com.agentclientprotocol.transport.BaseTransport
23+
import com.agentclientprotocol.transport.Transport
24+
import kotlinx.atomicfu.atomic
25+
import kotlinx.coroutines.CompletableDeferred
26+
import kotlinx.coroutines.CoroutineScope
27+
import kotlinx.coroutines.Dispatchers
28+
import kotlinx.coroutines.SupervisorJob
29+
import kotlinx.coroutines.async
30+
import kotlinx.coroutines.awaitAll
31+
import kotlinx.coroutines.cancel
32+
import kotlinx.coroutines.runBlocking
33+
import kotlinx.coroutines.withContext
34+
import kotlinx.coroutines.withTimeoutOrNull
35+
import kotlinx.serialization.json.JsonElement
36+
import kotlin.test.Test
37+
import kotlin.test.assertEquals
38+
import kotlin.time.Duration.Companion.seconds
39+
40+
/**
41+
* An agent may put a `session/update` on the wire immediately after the `session/new` response, so the client
42+
* receives both back to back while `newSession` is still returning on another thread. The notification must
43+
* still reach the session's operations: the client either queues it on the holder being initialized, or hands
44+
* it to the session that has just been registered.
45+
*
46+
* Getting this wrong used to be observable as `AcpExpectedError: Session <id> not found` logged by the
47+
* protocol, plus a silently dropped update, whenever the notification was handled while the concurrent
48+
* `newSession` was registering its holder and dropping `initializingSessionsCount` back to zero.
49+
*/
50+
class ClientSessionUpdateDeliveryTest {
51+
@Test
52+
fun `session update sent right after the session new response reaches the session`() {
53+
val delivered = atomic(0)
54+
// The interleaving is only reachable when the notification is handled while a concurrent newSession is
55+
// registering its session, so this is a probabilistic stress test: correct code can never fail it, while
56+
// the regression it guards against showed up in roughly one run out of three on a warm 10-core laptop.
57+
// Each batch gets a fresh client to keep the session map - and with it the cost of a round - from
58+
// growing over the run.
59+
repeat(BATCHES) {
60+
val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
61+
try {
62+
val transport = EagerUpdateAgentTransport()
63+
val protocol = Protocol(scope, transport)
64+
protocol.start()
65+
transport.start()
66+
val client = Client(protocol)
67+
68+
runBlocking {
69+
// Real parallelism between the protocol read loop and the `newSession` continuation is the
70+
// whole point, so sessions are created off the (single-threaded) runBlocking dispatcher.
71+
withContext(Dispatchers.Default) {
72+
repeat(ROUNDS_PER_BATCH) {
73+
(1..CONCURRENCY).map {
74+
async {
75+
val update = CompletableUpdate()
76+
client.newSession(SessionCreationParameters(cwd = ".", mcpServers = emptyList())) { _, _ ->
77+
update
78+
}
79+
if (update.awaitNotification()) delivered.incrementAndGet()
80+
}
81+
}.awaitAll()
82+
}
83+
}
84+
}
85+
} finally {
86+
scope.cancel()
87+
}
88+
}
89+
90+
assertEquals(BATCHES * ROUNDS_PER_BATCH * CONCURRENCY, delivered.value, "every session/update must reach its session")
91+
}
92+
93+
private companion object {
94+
private const val BATCHES = 60
95+
private const val ROUNDS_PER_BATCH = 250
96+
private const val CONCURRENCY = 4
97+
}
98+
}
99+
100+
/** Records the single `session/update` the fake agent sends for a session. */
101+
private class CompletableUpdate : ClientSessionOperations {
102+
private val notified = CompletableDeferred<Unit>()
103+
104+
override suspend fun requestPermissions(
105+
toolCall: SessionUpdate.ToolCallUpdate,
106+
permissions: List<PermissionOption>,
107+
_meta: JsonElement?,
108+
): RequestPermissionResponse = error("not expected in this test")
109+
110+
override suspend fun notify(notification: SessionUpdate, _meta: JsonElement?) {
111+
notified.complete(Unit)
112+
}
113+
114+
suspend fun awaitNotification(): Boolean = withTimeoutOrNull(NOTIFICATION_TIMEOUT) { notified.await() } != null
115+
116+
private companion object {
117+
private val NOTIFICATION_TIMEOUT = 5.seconds
118+
}
119+
}
120+
121+
/**
122+
* A minimal agent that answers `session/new` and, without yielding in between, sends one `session/update` for
123+
* the session it has just created - the sequence a real agent produces when it reports its available commands
124+
* as soon as the session exists.
125+
*/
126+
private class EagerUpdateAgentTransport : BaseTransport() {
127+
private val sessionCounter = atomic(0)
128+
129+
override fun start() {
130+
_state.value = Transport.State.STARTED
131+
}
132+
133+
override fun close() {
134+
_state.value = Transport.State.CLOSING
135+
fireClose()
136+
_state.value = Transport.State.CLOSED
137+
}
138+
139+
override fun send(message: JsonRpcMessage) {
140+
if (message !is JsonRpcRequest || message.method != AcpMethod.AgentMethods.SessionNew.methodName) return
141+
val sessionId = SessionId("session-${sessionCounter.incrementAndGet()}")
142+
fireMessage(
143+
JsonRpcResponse(
144+
id = message.id,
145+
result = ACPJson.encodeToJsonElement(
146+
AcpMethod.AgentMethods.SessionNew.responseSerializer,
147+
NewSessionResponse(sessionId),
148+
),
149+
)
150+
)
151+
fireMessage(
152+
JsonRpcNotification(
153+
method = AcpMethod.ClientMethods.SessionUpdate.methodName,
154+
params = ACPJson.encodeToJsonElement(
155+
AcpMethod.ClientMethods.SessionUpdate.serializer,
156+
SessionNotification(sessionId, SessionUpdate.AgentMessageChunk(ContentBlock.Text("update"))),
157+
),
158+
)
159+
)
160+
}
161+
}

0 commit comments

Comments
 (0)