#60 Add ability to schedule session post init activity

Author: nerzhulart

acp-ktor-test/src/commonTest/kotlin/com/agentclientprotocol/SimpleAgentTest.kt

@@ -970,4 +970,60 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
         assertEquals(1, sessions.size)
     }
 
+    @Test
+    @OptIn(UnstableApi::class)
+    fun `session ready notifies available commands`() = testWithProtocols { clientProtocol, agentProtocol ->
+        val readyUpdate = CompletableDeferred<SessionUpdate>()
+        val client = Client(protocol = clientProtocol)
+        Agent(protocol = agentProtocol, agentSupport = object : AgentSupport {
+            override suspend fun initialize(clientInfo: ClientInfo): AgentInfo {
+                return AgentInfo(clientInfo.protocolVersion)
+            }
+
+            override suspend fun createSession(sessionParameters: SessionCreationParameters): AgentSession {
+                return object : AgentSession {
+                    override val sessionId: SessionId = SessionId("ready-session-id")
+
+                    override suspend fun postInitialize() {
+                        currentCoroutineContext().client.notify(
+                            SessionUpdate.AvailableCommandsUpdate(
+                                listOf(AvailableCommand("help", "Show available commands", AvailableCommandInput.Unstructured("topic")))
+                            )
+                        )
+                    }
+
+                    override suspend fun prompt(
+                        content: List<ContentBlock>,
+                        _meta: JsonElement?,
+                    ): Flow<Event> = emptyFlow()
+                }
+            }
+        })
+
+        client.initialize(ClientInfo(protocolVersion = 10))
+        client.newSession(SessionCreationParameters("/test/path", emptyList())) { _, _ ->
+            object : ClientSessionOperations {
+                override suspend fun requestPermissions(
+                    toolCall: SessionUpdate.ToolCallUpdate,
+                    permissions: List<PermissionOption>,
+                    _meta: JsonElement?,
+                ): RequestPermissionResponse {
+                    return RequestPermissionResponse(RequestPermissionOutcome.Cancelled)
+                }
+
+                override suspend fun notify(
+                    notification: SessionUpdate,
+                    _meta: JsonElement?,
+                ) {
+                    readyUpdate.complete(notification)
+                }
+            }
+        }
+
+        val update = withTimeout(1000) { readyUpdate.await() }
+        assertTrue(update is SessionUpdate.AvailableCommandsUpdate)
+        val command = (update as SessionUpdate.AvailableCommandsUpdate).availableCommands.single()
+        assertEquals("help", command.name)
+    }
+
 }

acp/api/acp.api

@@ -49,6 +49,8 @@ public abstract interface class com/agentclientprotocol/agent/AgentSession {
 	public fun getDefaultMode-kyZWFqk ()Ljava/lang/String;
 	public fun getDefaultModel-GMZLII8 ()Ljava/lang/String;
 	public abstract fun getSessionId-7EW-EgU ()Ljava/lang/String;
+	public fun postInitialize (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun postInitialize$suspendImpl (Lcom/agentclientprotocol/agent/AgentSession;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public abstract fun prompt (Ljava/util/List;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun prompt$default (Lcom/agentclientprotocol/agent/AgentSession;Ljava/util/List;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public fun setConfigOption-p1da2xM (Ljava/lang/String;Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
@@ -66,6 +68,7 @@ public final class com/agentclientprotocol/agent/AgentSession$DefaultImpls {
 	public static fun getConfigOptions (Lcom/agentclientprotocol/agent/AgentSession;)Ljava/util/List;
 	public static fun getDefaultMode-kyZWFqk (Lcom/agentclientprotocol/agent/AgentSession;)Ljava/lang/String;
 	public static fun getDefaultModel-GMZLII8 (Lcom/agentclientprotocol/agent/AgentSession;)Ljava/lang/String;
+	public static fun postInitialize (Lcom/agentclientprotocol/agent/AgentSession;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun prompt$default (Lcom/agentclientprotocol/agent/AgentSession;Ljava/util/List;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static fun setConfigOption-p1da2xM (Lcom/agentclientprotocol/agent/AgentSession;Ljava/lang/String;Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static fun setMode-tFeI3nk (Lcom/agentclientprotocol/agent/AgentSession;Ljava/lang/String;Lkotlinx/serialization/json/JsonElement;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;

acp/src/commonMain/kotlin/com/agentclientprotocol/agent/Agent.kt

@@ -249,6 +249,7 @@ public class Agent(
             RemoteClientSessionOperations(protocol, session.sessionId, clientInfo.capabilities),
             protocol
         )
+        currentCoroutineContext().executeAfterCurrentRequest { sessionWrapper.executeWithSession { session.postInitialize() } }
 
         _sessions.update {
             it.put(session.sessionId, sessionWrapper)

acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt

@@ -9,6 +9,17 @@ import kotlinx.serialization.json.JsonElement
 public interface AgentSession {
     public val sessionId: SessionId
 
+    /**
+     * Executes after the session is created and send to the client. Can be used to send additional notifications like Commands.
+     *
+     * To access client operations use:
+     * ```
+     * currentCoroutineContext().client
+     * ```
+     *
+     * This method shouldn't throw exceptions.
+     */
+    public suspend fun postInitialize() {}
     /**
      * Sends a message to the agent for execution and waits for the whole turn to be completed.
      * During execution, the agent can send notifications or requests to the client.

acp/src/commonMain/kotlin/com/agentclientprotocol/protocol/Protocol.extensions.kt

@@ -111,11 +111,28 @@ public operator fun <TNotification : AcpNotification> AcpMethod.AcpNotificationM
     return rpc.sendNotification(this, notification)
 }
 
-internal class JsonRpcRequestContextElement(val jsonRpcRequest: JsonRpcRequest) : AbstractCoroutineContextElement(Key) {
+internal class RequestHolder(val jsonRpcRequest: JsonRpcRequest) {
+    // probably make it thread safe
+    internal val handlers = mutableListOf<suspend () -> Unit>()
+    fun executeAfterCurrentRequest(block: suspend () -> Unit) {
+        handlers.add(block)
+    }
+}
+
+internal class JsonRpcRequestContextElement(val requestHolder: RequestHolder) : AbstractCoroutineContextElement(Key) {
     object Key : CoroutineContext.Key<JsonRpcRequestContextElement>
 }
 
+internal val CoroutineContext.requestHolder: RequestHolder
+    get() = this[JsonRpcRequestContextElement.Key]?.requestHolder ?: error("There is no active incoming request in this context")
+
 public val CoroutineContext.jsonRpcRequest: JsonRpcRequest
-    get() = this[JsonRpcRequestContextElement.Key]?.jsonRpcRequest ?: error("No JsonRpcRequest found in context")
+    get() = this.requestHolder.jsonRpcRequest
 
-internal fun JsonRpcRequest.asContextElement() = JsonRpcRequestContextElement(this)
+
+/**
+ * Execute a block after the current request is processed and the response is sent back to the client.
+ */
+internal fun CoroutineContext.executeAfterCurrentRequest(block: suspend () -> Unit) {
+    requestHolder.executeAfterCurrentRequest(block)
+}

acp/src/commonMain/kotlin/com/agentclientprotocol/protocol/Protocol.kt

@@ -372,10 +372,11 @@ public class Protocol(
     }
 
     private suspend fun handleRequest(request: JsonRpcRequest) {
+        val requestHolder = RequestHolder(request)
         val handler = requestHandlers.value[request.method]
         if (handler != null) {
             try {
-                val result = withContext(request.asContextElement()) {
+                val result = withContext(JsonRpcRequestContextElement(requestHolder)) {
                     handler(request)
                 }
                 sendResponse(request.id, result, null)
@@ -415,6 +416,16 @@ public class Protocol(
                     )
                 )
             }
+            for (handler in requestHolder.handlers) {
+                runCatching { handler() }.onFailure { t ->
+                    if (t is CancellationException) {
+                        // ignore CE
+                        logger.trace(t) { "Request handler for '${request.method}' cancelled" }
+                    } else {
+                        logger.error(t) { "Error handling after request handlers for ${request.method}" }
+                    }
+                }
+            }
         } else {
             val error = JsonRpcError(
                 code = JsonRpcErrorCode.METHOD_NOT_FOUND.code, message = "Method not supported: ${request.method}"

build.gradle.kts

@@ -7,7 +7,7 @@ plugins {
 private val buildNumber: String? = System.getenv("GITHUB_RUN_NUMBER")
 private val isReleasePublication = System.getenv("RELEASE_PUBLICATION")?.toBoolean() ?: false
 
-private val baseVersion = "0.14.1"
+private val baseVersion = "0.15.1"
 
 allprojects {
     group = "com.agentclientprotocol"