[IJAI-1133]: fix(client): don't fail session/update for sessions the client never connected to

acp/api/acp.api

@@ -276,8 +276,8 @@ public final class com/agentclientprotocol/agent/v2/SessionCreationParameters {
 }
 
 public final class com/agentclientprotocol/client/Client {
-	public fun <init> (Lcom/agentclientprotocol/protocol/Protocol;Lcom/agentclientprotocol/client/GlobalElicitationHandler;)V
-	public synthetic fun <init> (Lcom/agentclientprotocol/protocol/Protocol;Lcom/agentclientprotocol/client/GlobalElicitationHandler;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lcom/agentclientprotocol/protocol/Protocol;Lcom/agentclientprotocol/client/GlobalElicitationHandler;Lcom/agentclientprotocol/client/GlobalSessionUpdateHandler;)V
+	public synthetic fun <init> (Lcom/agentclientprotocol/protocol/Protocol;Lcom/agentclientprotocol/client/GlobalElicitationHandler;Lcom/agentclientprotocol/client/GlobalSessionUpdateHandler;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun authenticate-fMnwWJU (Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun authenticate-fMnwWJU$default (Lcom/agentclientprotocol/client/Client;Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public final fun deleteSession-nk3TnMc (Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
@@ -288,6 +288,7 @@ public final class com/agentclientprotocol/client/Client {
 	public final fun getAgentInfo ()Lcom/agentclientprotocol/agent/AgentInfo;
 	public final fun getClientInfo ()Lcom/agentclientprotocol/client/ClientInfo;
 	public final fun getGlobalElicitationHandler ()Lcom/agentclientprotocol/client/GlobalElicitationHandler;
+	public final fun getGlobalSessionUpdateHandler ()Lcom/agentclientprotocol/client/GlobalSessionUpdateHandler;
 	public final fun getNesSession-0izbxq0 (Ljava/lang/String;)Lcom/agentclientprotocol/client/ClientNesSession;
 	public final fun getProtocol ()Lcom/agentclientprotocol/protocol/Protocol;
 	public final fun getSession-0izbxq0 (Ljava/lang/String;)Lcom/agentclientprotocol/client/ClientSession;
@@ -426,6 +427,10 @@ public abstract interface class com/agentclientprotocol/client/GlobalElicitation
 	public abstract fun createElicitation (Lcom/agentclientprotocol/model/CreateElicitationRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }
 
+public abstract interface class com/agentclientprotocol/client/GlobalSessionUpdateHandler {
+	public abstract fun onUnconnectedSessionUpdate-wPMwmcM (Ljava/lang/String;Lcom/agentclientprotocol/model/SessionUpdate;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public abstract interface class com/agentclientprotocol/client/NegotiatedClient {
 	public abstract fun getProtocolVersion ()I
 }

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

@@ -43,7 +43,9 @@ public typealias ClientInstance = Client
 public class Client(
     public val protocol: Protocol,
     @property:UnstableApi
-    public val globalElicitationHandler: GlobalElicitationHandler? = null
+    public val globalElicitationHandler: GlobalElicitationHandler? = null,
+    @property:UnstableApi
+    public val globalSessionUpdateHandler: GlobalSessionUpdateHandler? = null
 ) {
     private class ClientSessionHolder {
         private val sessionDeferred: CompletableDeferred<ClientSessionImpl> = CompletableDeferred()
@@ -72,6 +74,20 @@ public class Client(
             sessionDeferred.completeExceptionally(cause)
         }
 
+        /**
+         * Like [completeExceptionally], but also returns whatever was queued before this call - for a holder
+         * that was speculatively created for a session that never actually got claimed (see the
+         * `initializingSessionsCount > 0` branch of [findSessionHolder]), those notifications belong to no live
+         * session and would otherwise be silently discarded.
+         */
+        fun completeExceptionallyAndDrainQueue(cause: Throwable): List<Pair<SessionUpdate, JsonElement?>> {
+            val drained = buildList {
+                while (true) add(notifications.tryReceive().getOrNull() ?: break)
+            }
+            completeExceptionally(cause)
+            return drained
+        }
+
         suspend fun handleOrQueue(notification: SessionUpdate, _meta: JsonElement?) {
             val sendResult = notifications.trySend(Pair(notification, _meta))
 
@@ -94,9 +110,12 @@ public class Client(
     private val _elicitationToSession = ElicitationSessionStore()
 
     /**
-     * Creates a new entry only if there are some currently initializing sessions. Otherwise, throws in the case of missing session.
+     * Looks up the holder for [sessionId], creating a new entry only if there are some currently initializing
+     * sessions. Returns `null` if the session is neither registered nor being initialized, instead of throwing -
+     * callers that need a session to exist (e.g., to service a request against it) should use
+     * [getOrCreateSessionHolder] instead.
      */
-    private fun getOrCreateSessionHolder(sessionId: SessionId): ClientSessionHolder {
+    private fun findSessionHolder(sessionId: SessionId): ClientSessionHolder? {
         // Fast path for the common case of an already registered session.
         _sessions.value.sessions[sessionId]?.let { return it }
         var clientSessionHolder: ClientSessionHolder? = null
@@ -122,9 +141,15 @@ public class Client(
                 }
             }
         }
-        return clientSessionHolder ?: acpFail("Session $sessionId not found")
+        return clientSessionHolder
     }
 
+    /**
+     * Creates a new entry only if there are some currently initializing sessions. Otherwise, throws in the case of missing session.
+     */
+    private fun getOrCreateSessionHolder(sessionId: SessionId): ClientSessionHolder =
+        findSessionHolder(sessionId) ?: acpFail("Session $sessionId not found")
+
     internal fun removeSessionHolder(sessionId: SessionId) {
         _sessions.update { currentMap ->
             currentMap.copy(sessions = currentMap.sessions.remove(sessionId))
@@ -214,7 +239,20 @@ public class Client(
         }
 
         protocol.setNotificationHandler(AcpMethod.ClientMethods.V1.SessionUpdate) { params: SessionNotification ->
-            val sessionHolder = getOrCreateSessionHolder(params.sessionId)
+            // The agent may report an update (e.g., a status change on `session/list`) for a session this client
+            // never called `session/new` / `session/load` / `session/resume` for. It can live on the server,
+            // created from another IDE window, the web, or another machine. That's not a protocol violation, so
+            // unlike other session-scoped methods, an unknown/unconnected session here must not fail the call.
+            val sessionHolder = findSessionHolder(params.sessionId)
+            if (sessionHolder == null) {
+                val handler = globalSessionUpdateHandler
+                if (handler != null) {
+                    handler.onUnconnectedSessionUpdate(params.sessionId, params.update, params._meta)
+                } else {
+                    logger.debug { "Ignoring session/update for session ${params.sessionId}: client is not connected to it" }
+                }
+                return@setNotificationHandler
+            }
             sessionHolder.handleOrQueue(params.update, params._meta)
         }
 
@@ -684,8 +722,17 @@ public class Client(
             if (hangingSessions != null) {
                 for ((id, holder) in hangingSessions) {
                     logger.trace { "Removing hanging session $id" }
-                    // report it as non existent session
-                    holder.completeExceptionally(AcpExpectedError("Session $id not found"))
+                    // report it as a non-existent session
+                    val queuedUpdates = holder.completeExceptionallyAndDrainQueue(AcpExpectedError("Session $id not found"))
+                    // These were buffered on the assumption they might belong to this (or another concurrent)
+                    // initialization; since none claimed `id`, it's an unconnected session, same as if no
+                    // initialization had been in progress when its updates arrived.
+                    val handler = globalSessionUpdateHandler
+                    if (handler != null) {
+                        for ((update, meta) in queuedUpdates) {
+                            handler.onUnconnectedSessionUpdate(id, update, meta)
+                        }
+                    }
                 }
             }
         }

acp/src/commonMain/kotlin/com/agentclientprotocol/client/GlobalSessionUpdateHandler.kt

@@ -0,0 +1,19 @@
+package com.agentclientprotocol.client
+
+import com.agentclientprotocol.annotations.UnstableApi
+import com.agentclientprotocol.model.SessionId
+import com.agentclientprotocol.model.SessionUpdate
+import kotlinx.serialization.json.JsonElement
+
+/**
+ * Handler for `session/update` notifications about a session the client is not connected to.
+ *
+ * A session can live on the server without this client ever having called `session/new` / `session/load` /
+ * `session/resume` for it - e.g. it was created from another IDE window, the web, or another machine. This is
+ * invoked instead of failing or silently dropping the notification, letting a client observe such updates -
+ * for example to keep a `session/list`-rendered list live without polling.
+ */
+@UnstableApi
+public fun interface GlobalSessionUpdateHandler {
+    public suspend fun onUnconnectedSessionUpdate(sessionId: SessionId, update: SessionUpdate, _meta: JsonElement?)
+}

acp/src/jvmTest/kotlin/com/agentclientprotocol/client/ClientGlobalSessionUpdateHandlerTest.kt

@@ -0,0 +1,85 @@
+@file:OptIn(UnstableApi::class)
+
+package com.agentclientprotocol.client
+
+import com.agentclientprotocol.annotations.UnstableApi
+import com.agentclientprotocol.model.*
+import com.agentclientprotocol.protocol.Protocol
+import com.agentclientprotocol.rpc.ACPJson
+import com.agentclientprotocol.rpc.JsonRpcMessage
+import com.agentclientprotocol.rpc.JsonRpcNotification
+import com.agentclientprotocol.transport.BaseTransport
+import com.agentclientprotocol.transport.Transport
+import kotlinx.coroutines.*
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+import kotlin.time.Duration.Companion.seconds
+
+/**
+ * A session can live on the server without this client ever having called `session/new` / `session/load` /
+ * `session/resume` for it. [GlobalSessionUpdateHandler] lets a client observe such updates - e.g. to keep a
+ * `session/list`-rendered list live without polling - instead of the update being silently dropped.
+ *
+ * See https://youtrack.jetbrains.com/issue/IJAI-1133
+ */
+class ClientGlobalSessionUpdateHandlerTest {
+    @Test
+    fun `session update for an unconnected session is delivered to the global session update handler`() {
+        val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
+        try {
+            val transport = NotifyingTransport()
+            val protocol = Protocol(scope, transport)
+            protocol.start()
+            transport.start()
+
+            val received = CompletableDeferred<Pair<SessionId, SessionUpdate>>()
+            Client(
+                protocol,
+                globalSessionUpdateHandler = { sessionId, update, _ ->
+                    received.complete(sessionId to update)
+                },
+            )
+
+            val sessionId = SessionId("unconnected-session")
+            val update = SessionUpdate.AgentMessageChunk(ContentBlock.Text("update"))
+
+            runBlocking {
+                transport.emitSessionUpdate(sessionId, update)
+                val result = withTimeoutOrNull(5.seconds) { received.await() }
+                assertNotNull(result, "the global session update handler must be invoked for an unconnected session")
+                assertEquals(sessionId, result.first)
+                assertEquals(update, result.second)
+            }
+        } finally {
+            scope.cancel()
+        }
+    }
+}
+
+/** A transport whose only job is to let the test push arbitrary `session/update` notifications to the client. */
+private class NotifyingTransport : BaseTransport() {
+    override fun start() {
+        _state.value = Transport.State.STARTED
+    }
+
+    override fun close() {
+        _state.value = Transport.State.CLOSING
+        fireClose()
+        _state.value = Transport.State.CLOSED
+    }
+
+    override fun send(message: JsonRpcMessage) = Unit
+
+    fun emitSessionUpdate(sessionId: SessionId, update: SessionUpdate) {
+        fireMessage(
+            JsonRpcNotification(
+                method = AcpMethod.ClientMethods.SessionUpdate.methodName,
+                params = ACPJson.encodeToJsonElement(
+                    AcpMethod.ClientMethods.SessionUpdate.serializer,
+                    SessionNotification(sessionId, update),
+                ),
+            )
+        )
+    }
+}