diff --git a/AGENTS.md b/AGENTS.md index f76c93a..180dcb1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -149,6 +149,7 @@ Lifecycle correctness is part of the public behavior. Changes to startup, cancel - Return follow-up work and cleanup through `RequestOutcome` rather than running them inside the handler. - Keep normal close, timeout/abort, parent cancellation, and I/O failure paths distinct and testable. - The v1 client keeps one `ClientSessionHolder` per session id. With a `ClientOperationsFactory`, the holder buffers the session's updates until the opening call answers. The `loadSession` and `resumeSession` overloads that take `ClientSessionOperations` claim the holder before the request goes out, and one consumer hands its queue to the operations until the answer. Until the switch to the session, every update passes that one queue with its one consumer, so the arrival order holds by construction: do not add a second delivery path, a lock, or a counter there. After the switch, an update goes synchronously to the session. An update that misses the closed queue waits in the `lateDeliveries` chain for the drain and for the update that missed the queue before it, so late updates keep their order too. +- The v1 agent runs one turn at a time per session, an SDK guarantee the protocol does not define. Only the turn frees its slot, when it answers, so a `session/prompt` sent earlier, even while a cancelled turn winds down, is rejected. `session/cancel` targets the turn registered when it arrives, never one registered while `AgentSession.cancel()` runs, and a cancelled turn answers only after every such call returns, so updates sent from it precede the response. A turn's flow failure stays inside the turn's job so it still passes that wait: after a cancel it answers `cancelled` (aborted calls often throw), otherwise it answers with the error. For WebSocket shutdown, preserve the established sequence: reject new work, enter closing state, drain accepted frames within the configured timeout, then perform a normal WebSocket close. Reserve cancellation/abort for timeout, cancellation, or writer failure. diff --git a/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/Agent.kt b/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/Agent.kt index 033e0ca..bf0bad1 100644 --- a/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/Agent.kt +++ b/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/Agent.kt @@ -12,11 +12,16 @@ import io.github.oshai.kotlinlogging.KotlinLogging import kotlinx.atomicfu.atomic import kotlinx.atomicfu.update import kotlinx.collections.immutable.persistentMapOf +import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.Job import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.first +import kotlinx.coroutines.flow.getAndUpdate +import kotlinx.coroutines.flow.update import kotlinx.coroutines.launch import kotlinx.coroutines.withContext import kotlinx.serialization.json.JsonElement @@ -62,14 +67,60 @@ public class Agent( val clientOperations: ClientSessionOperations, protocol: Protocol ) : BaseSessionWrapper(agent, protocol) { - private class PromptSession(val promptJob: Job) + private class PromptSession(val promptJob: Job) { + /** + * The number of cancels still running for this turn, or [ANSWERING] once the turn has passed the wait for + * them and is about to answer. + */ + private val runningCancels = MutableStateFlow(0) + + private val _isCancelRequested = atomic(false) + + /** + * Whether a `session/cancel` targeted this turn before it started answering. + */ + val isCancelRequested: Boolean get() = _isCancelRequested.value + + /** + * Registers a cancel for this turn. Returns `false` if the turn is already answering. + */ + fun startCancel(): Boolean { + val started = runningCancels.getAndUpdate { if (it == ANSWERING) it else it + 1 } != ANSWERING + if (started) _isCancelRequested.value = true + return started + } + + fun finishCancel() { + runningCancels.update { it - 1 } + } + + /** + * Suspends until no cancel is running for this turn, then marks the turn as answering. + */ + suspend fun awaitCancels() { + while (!runningCancels.compareAndSet(0, ANSWERING)) { + runningCancels.first { it == 0 } + } + } + + private companion object { + const val ANSWERING = -1 + } + } + private val _activePrompt = atomic(null) + @OptIn(UnstableApi::class) suspend fun prompt(content: List, _meta: JsonElement? = null): PromptResponse { var response: PromptResponse? = null + var failure: Throwable? = null return coroutineScope { - try { - val promptJob = launch(start = CoroutineStart.LAZY) { + val promptJob = launch(start = CoroutineStart.LAZY) { + /* + * Keep a failure inside the turn's job, so a failed turn still waits for running cancels and a + * cancel can still decide its answer. + */ + try { agentSession.prompt(content, _meta).collect { event -> when (event) { is Event.PromptResponseEvent -> { @@ -84,37 +135,66 @@ public class Agent( } } } + } catch (e: CancellationException) { + throw e + } catch (e: Throwable) { + failure = e } + } - val promptSession = PromptSession(promptJob) - if (!_activePrompt.compareAndSet(null, promptSession)) { - error("There is already active prompt execution") - } + val promptSession = PromptSession(promptJob) + if (!_activePrompt.compareAndSet(null, promptSession)) { + error("There is already active prompt execution") + } + try { promptJob.join() + /* + * A cancelled turn answers only after the session's cancel() returns: the response must follow + * every aborted operation and every update, including those sent from cancel(). + */ + promptSession.awaitCancels() + /* + * The session's cancel() may stop the turn cooperatively, so its flow ends normally before the + * SDK cancels the job. The client still cancelled the turn and must get the `cancelled` reason. + */ + val isCancelled = promptJob.isCancelled || promptSession.isCancelRequested + failure?.let { e -> + if (!isCancelled) throw e + logger.debug(e) { "Prompt turn failed after a cancel; answering as cancelled" } + } response ?: PromptResponse( - stopReason = if (promptJob.isCancelled) StopReason.CANCELLED else StopReason.END_TURN + stopReason = if (isCancelled) StopReason.CANCELLED else StopReason.END_TURN ) } finally { - _activePrompt.getAndSet(null) + // Only the turn frees its slot, so the next prompt can start only after this one has answered. + _activePrompt.value = null } } } suspend fun cancel() { - // notify AgentSession about upcoming cancellation, this way implementations can gracefully stop ongoing requests - agentSession.cancel() - - val activePrompt = _activePrompt.getAndSet(null) - if (activePrompt != null) { - logger.trace { "Cancelling prompt" } - // we expect that all nested outgoing jobs will be cancelled automatically due to structured concurrency - // -> prompt task - // <- [request] read file - // -> [response] read file - // <- [request] permissions - // |suspended| - // cancelling the whole prompt should cancel all nested outgoing requests. These requests on CE will propagate cancellation to the other side - activePrompt.promptJob.cancel() + /* + * Take the turn to cancel before notifying the session: the session's cancel() may suspend, and a turn + * that has already answered may be followed by the next one meanwhile, which this cancel must not touch. + */ + val promptToCancel = _activePrompt.value?.takeIf { it.startCancel() } + try { + // notify AgentSession about upcoming cancellation, this way implementations can gracefully stop ongoing requests + agentSession.cancel() + + if (promptToCancel != null) { + logger.trace { "Cancelling prompt" } + // we expect that all nested outgoing jobs will be cancelled automatically due to structured concurrency + // -> prompt task + // <- [request] read file + // -> [response] read file + // <- [request] permissions + // |suspended| + // cancelling the whole prompt should cancel all nested outgoing requests. These requests on CE will propagate cancellation to the other side + promptToCancel.promptJob.cancel() + } + } finally { + promptToCancel?.finishCancel() } } } diff --git a/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt b/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt index 7e84b87..b129b52 100644 --- a/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt +++ b/acp/src/commonMain/kotlin/com/agentclientprotocol/agent/AgentSession.kt @@ -41,6 +41,11 @@ public interface AgentSession { * 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. * + * The SDK runs at most one turn per session: this method is not called again until the previous turn's + * `session/prompt` has been answered, including a cancelled turn that is still winding down. A `session/prompt` + * that arrives earlier is rejected with an internal error (`-32603`). The protocol does not define overlapping + * prompts; this is an SDK guarantee. + * * Corresponds to the [AcpMethod.AgentMethods.V1.SessionPrompt] */ public suspend fun prompt(content: List, _meta: JsonElement? = null): Flow @@ -48,6 +53,19 @@ public interface AgentSession { /** * Cancels the current agent turn and returns after the agent canceled all activities of the current turn. * + * The cancelled turn's `session/prompt` is answered only after this method returns, even if the [prompt] flow + * ends or throws earlier, so updates sent from here through `currentCoroutineContext().client` reach the client + * before the response. This method must therefore not wait for that response. A turn that starts after this call + * began is not cancelled by it. + * + * The cancelled turn is answered with [StopReason.CANCELLED] unless its [prompt] flow emitted a + * [Event.PromptResponseEvent], even if this method stops the flow cooperatively and it completes normally. An + * exception the flow throws after the cancel, such as an aborted call's error, does not fail the turn: it is + * logged at debug level and the turn is answered the same way. A turn that fails without a cancel is answered + * with the error. + * + * Each `session/cancel` calls this method, so a call may start while a previous one is still running. + * * Corresponds to the [AcpMethod.AgentMethods.V1.SessionCancel] */ public suspend fun cancel() {} diff --git a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/AgentTest.kt b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/AgentTest.kt index b069843..dbc12fc 100644 --- a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/AgentTest.kt +++ b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/AgentTest.kt @@ -1,17 +1,20 @@ package com.agentclientprotocol.agent import com.agentclientprotocol.annotations.UnstableApi +import com.agentclientprotocol.common.Event import com.agentclientprotocol.common.SessionCreationParameters import com.agentclientprotocol.model.* import com.agentclientprotocol.protocol.Protocol import com.agentclientprotocol.protocol.setNotificationHandler import com.agentclientprotocol.rpc.* import kotlinx.coroutines.* +import kotlinx.coroutines.channels.Channel import kotlinx.coroutines.flow.emptyFlow import kotlinx.coroutines.flow.flow import kotlinx.serialization.Serializable import kotlinx.serialization.json.* import kotlin.test.* +import kotlin.time.Duration.Companion.milliseconds import kotlin.time.Duration.Companion.seconds import com.agentclientprotocol.agent.v2.Agent as V2Agent import com.agentclientprotocol.agent.v2.AgentInfo as V2AgentInfo @@ -933,6 +936,324 @@ class AgentTest { assertEquals(StopReason.CANCELLED, response.stopReason) } } + + /* + * The SDK runs one turn at a time per session, and a cancelled turn still runs until it answers. A prompt sent + * while the cancelled turn winds down is rejected, and the next prompt after the answer runs and can be + * cancelled. + */ + @Test + fun `a prompt sent before the cancelled turn answers is rejected`() { + val started = Channel(Channel.UNLIMITED) + val firstCancelled = CompletableDeferred() + val releaseFirst = CompletableDeferred() + val promptHandler: PromptHandler = { content -> + val text = assertIs(content.single()).text + started.send(text) + try { + awaitCancellation() + } catch (e: CancellationException) { + if (text == "first") { + // Holds the cancelled turn open after the SDK's cancel() has returned. + firstCancelled.complete(Unit) + withContext(NonCancellable) { releaseFirst.await() } + } + throw e + } + } + withTestAgentSession(promptHandler = promptHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(101)) + assertEquals("first", withTimeout(5.seconds) { started.receive() }) + testAgent.testCancel(CancelNotification(session.sessionId)) + withTimeout(5.seconds) { firstCancelled.await() } + + try { + testAgent.firePrompt(session.sessionId, "early", RequestId.create(102)) + testAgent.receivePromptRejection(RequestId.create(102)) + } finally { + releaseFirst.complete(Unit) + } + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(101)).stopReason) + + testAgent.firePrompt(session.sessionId, "next", RequestId.create(103)) + assertEquals("next", withTimeout(5.seconds) { started.receive() }) + testAgent.testCancel(CancelNotification(session.sessionId)) + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(103)).stopReason) + } + } + + @Test + fun `a rejected overlapping prompt leaves the running turn cancellable`() { + val started = CompletableDeferred() + val promptHandler: PromptHandler = { + started.complete(Unit) + awaitCancellation() + } + withTestAgentSession(promptHandler = promptHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(201)) + withTimeout(5.seconds) { started.await() } + + testAgent.firePrompt(session.sessionId, "overlapping", RequestId.create(202)) + testAgent.receivePromptRejection(RequestId.create(202)) + + testAgent.testCancel(CancelNotification(session.sessionId)) + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(201)).stopReason) + } + } + + /* + * AgentSession.cancel() returns after the turn's activities are cancelled, and the turn's flow may end before + * that. The agent MUST answer the cancelled prompt only after all ongoing operations are aborted and pending + * updates are sent, so updates sent from the session's cancel() precede the response. + * https://agentclientprotocol.com/protocol/prompt-turn#cancellation + */ + @Test + fun `a cancelled turn answers only after the session's cancel returns`() { + val started = CompletableDeferred() + val stopTurn = CompletableDeferred() + val turnEnded = CompletableDeferred() + val releaseCancel = CompletableDeferred() + + val promptHandler: PromptHandler = { + started.complete(Unit) + try { + stopTurn.await() + emit(Event.PromptResponseEvent(PromptResponse(StopReason.CANCELLED))) + } finally { + turnEnded.complete(Unit) + } + } + + val cancelHandler: CancelHandler = { + stopTurn.complete(Unit) + releaseCancel.await() + currentCoroutineContext().client.notify(SessionUpdate.AgentMessageChunk(ContentBlock.Text("cleanup"))) + } + + withTestAgentSession(promptHandler = promptHandler, cancelHandler = cancelHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(301)) + withTimeout(5.seconds) { started.await() } + testAgent.testCancel(CancelNotification(session.sessionId)) + withTimeout(5.seconds) { turnEnded.await() } + + /* + * The absence of a reply can only be observed for a while: the turn has ended, so without the fix + * the reply is sent right away. + */ + val early = withTimeoutOrNull(300.milliseconds) { testAgent.transport.receiveTestMessages(1) } + assertNull(early, "The cancelled turn answered before the session's cancel returned") + + releaseCancel.complete(Unit) + val update = assertIs(testAgent.transport.receiveTestMessages(1).singleOrNull()) + assertEquals(AcpMethod.ClientMethods.V1.SessionUpdate.methodName, update.method) + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(301)).stopReason) + } + } + + /* + * The session's cancel() may stop the turn cooperatively, so its flow ends normally before the SDK cancels the + * job. The client cancelled the turn, so the agent MUST answer with the `cancelled` stop reason. + * https://agentclientprotocol.com/protocol/prompt-turn#cancellation + */ + @Test + fun `a turn stopped cooperatively by the session's cancel answers cancelled`() { + val started = CompletableDeferred() + val stopTurn = CompletableDeferred() + val turnEnded = CompletableDeferred() + + val promptHandler: PromptHandler = { + started.complete(Unit) + stopTurn.await() + turnEnded.complete(Unit) + } + + val cancelHandler: CancelHandler = { + stopTurn.complete(Unit) + turnEnded.await() + } + + withTestAgentSession(promptHandler = promptHandler, cancelHandler = cancelHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(321)) + withTimeout(5.seconds) { started.await() } + testAgent.testCancel(CancelNotification(session.sessionId)) + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(321)).stopReason) + } + } + + /* + * Control for the cooperative stop: a response the turn emits itself is kept, even after a cancel. + */ + @Test + fun `a turn stopped cooperatively keeps the response it emitted`() { + val started = CompletableDeferred() + val stopTurn = CompletableDeferred() + val turnEnded = CompletableDeferred() + + val promptHandler: PromptHandler = { + started.complete(Unit) + stopTurn.await() + emit(Event.PromptResponseEvent(PromptResponse(StopReason.MAX_TOKENS))) + turnEnded.complete(Unit) + } + + val cancelHandler: CancelHandler = { + stopTurn.complete(Unit) + turnEnded.await() + } + + withTestAgentSession(promptHandler = promptHandler, cancelHandler = cancelHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(331)) + withTimeout(5.seconds) { started.await() } + testAgent.testCancel(CancelNotification(session.sessionId)) + assertEquals(StopReason.MAX_TOKENS, testAgent.receivePromptResponse(RequestId.create(331)).stopReason) + } + } + + /* + * Aborting an operation from the session's cancel() often makes it throw, for example an HTTP client's + * IOException("Canceled"). The agent MUST catch such errors and answer with the `cancelled` stop reason, and + * still only after the session's cancel() returns. + * https://agentclientprotocol.com/protocol/prompt-turn#cancellation + */ + @Test + fun `a turn that fails after the session's cancel answers cancelled after the cancel returns`() { + val started = CompletableDeferred() + val stopTurn = CompletableDeferred() + val turnEnded = CompletableDeferred() + val releaseCancel = CompletableDeferred() + + val promptHandler: PromptHandler = { + started.complete(Unit) + try { + stopTurn.await() + throw IllegalStateException("Canceled") + } finally { + turnEnded.complete(Unit) + } + } + + val cancelHandler: CancelHandler = { + stopTurn.complete(Unit) + releaseCancel.await() + currentCoroutineContext().client.notify(SessionUpdate.AgentMessageChunk(ContentBlock.Text("cleanup"))) + } + + withTestAgentSession(promptHandler = promptHandler, cancelHandler = cancelHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(341)) + withTimeout(5.seconds) { started.await() } + testAgent.testCancel(CancelNotification(session.sessionId)) + withTimeout(5.seconds) { turnEnded.await() } + + /* + * The absence of a reply can only be observed for a while: the turn has failed, so without the fix + * the error reply is sent right away. + */ + val early = withTimeoutOrNull(300.milliseconds) { testAgent.transport.receiveTestMessages(1) } + assertNull(early, "The failed cancelled turn answered before the session's cancel returned") + + releaseCancel.complete(Unit) + val update = assertIs(testAgent.transport.receiveTestMessages(1).singleOrNull()) + assertEquals(AcpMethod.ClientMethods.V1.SessionUpdate.methodName, update.method) + assertEquals(StopReason.CANCELLED, testAgent.receivePromptResponse(RequestId.create(341)).stopReason) + } + } + + /* + * Control for the failed cancelled turn: a turn that fails without a cancel still answers with the error. + */ + @Test + fun `a turn that fails without a cancel answers with the error`() { + val promptHandler: PromptHandler = { + throw IllegalStateException("Turn failed") + } + + withTestAgentSession(promptHandler = promptHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(351)) + val message = assertNotNull(testAgent.transport.receiveTestMessages(1).singleOrNull(), "No reply to the prompt") + val error = assertIs(message) + assertEquals(RequestId.create(351), error.id) + assertEquals(JsonRpcErrorCode.INTERNAL_ERROR.code, error.error.code) + assertEquals("Turn failed", error.error.message) + } + } + + /* + * A session/cancel can arrive just after the turn has answered, and AgentSession.cancel() may still suspend. + * The client may start the next turn meanwhile, and that turn was never asked to stop: it must not be + * cancelled. + */ + @Test + fun `a cancel that arrives after the turn answered leaves the next prompt running`() { + val started = Channel(Channel.UNLIMITED) + val cancelStarted = CompletableDeferred() + val releaseCancel = CompletableDeferred() + val cancelReturned = CompletableDeferred() + + val promptHandler: PromptHandler = { content -> + val text = assertIs(content.single()).text + started.send(text) + if (text == "second") { + /* + * Completed by the session's cancel() on the single handler thread, so this turn resumes only + * after the SDK's cancel() has finished with the active-prompt slot. + */ + cancelReturned.await() + } + emit(Event.PromptResponseEvent(PromptResponse(StopReason.END_TURN))) + } + + val cancelHandler: CancelHandler = { + cancelStarted.complete(Unit) + releaseCancel.await() + cancelReturned.complete(Unit) + } + + withTestAgentSession(promptHandler = promptHandler, cancelHandler = cancelHandler) { testAgent, session -> + testAgent.firePrompt(session.sessionId, "first", RequestId.create(311)) + assertEquals(StopReason.END_TURN, testAgent.receivePromptResponse(RequestId.create(311)).stopReason) + testAgent.testCancel(CancelNotification(session.sessionId)) + withTimeout(5.seconds) { cancelStarted.await() } + + testAgent.firePrompt(session.sessionId, "second", RequestId.create(312)) + assertEquals("first", withTimeout(5.seconds) { started.receive() }) + assertEquals("second", withTimeout(5.seconds) { started.receive() }) + releaseCancel.complete(Unit) + assertEquals(StopReason.END_TURN, testAgent.receivePromptResponse(RequestId.create(312)).stopReason) + } + } + + /** + * Sends a `session/prompt` without waiting for its reply, so that several turns can be in flight. + */ + private fun TestAgent.firePrompt(sessionId: SessionId, text: String, id: RequestId) { + val prompt = AcpMethod.AgentMethods.V1.SessionPrompt + transport.fireTestRequestWithoutReply( + prompt.methodName, + ACPJson.encodeToJsonElement(prompt.requestSerializer, PromptRequest(sessionId, listOf(ContentBlock.Text(text)))), + id, + ) + } + + /** + * Reads the next message from the agent as the rejection of the prompt [id] that overlapped a running turn. + */ + private suspend fun TestAgent.receivePromptRejection(id: RequestId) { + val message = assertNotNull(transport.receiveTestMessages(1).singleOrNull(), "No reply to prompt $id") + val rejection = assertIs(message) + assertEquals(id, rejection.id) + assertEquals(JsonRpcErrorCode.INTERNAL_ERROR.code, rejection.error.code) + assertEquals("There is already active prompt execution", rejection.error.message) + } + + /** + * Reads the next message from the agent as the successful reply to the prompt [id]. + */ + private suspend fun TestAgent.receivePromptResponse(id: RequestId): PromptResponse { + val message = assertNotNull(transport.receiveTestMessages(1).singleOrNull(), "No reply to prompt $id") + val reply = assertIs(message) + assertEquals(id, reply.id) + return ACPJson.decodeFromJsonElement(AcpMethod.AgentMethods.V1.SessionPrompt.responseSerializer, reply.result) + } } @OptIn(UnstableApi::class) diff --git a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestAgent.kt b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestAgent.kt index 13d45c3..aca8588 100644 --- a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestAgent.kt +++ b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestAgent.kt @@ -99,6 +99,7 @@ suspend fun TestAgent.simplePrompt(prompt: String): Pair.(List) -> Unit +typealias CancelHandler = suspend () -> Unit class TestAgentSession( override val sessionId: SessionId, - val promptHandler: PromptHandler + val promptHandler: PromptHandler, + val cancelHandler: CancelHandler = {}, ) : AgentSession { override suspend fun prompt(content: List, _meta: JsonElement?): Flow = flow { promptHandler(content) } + + override suspend fun cancel() { + cancelHandler() + } } fun withTestAgent( timeout: Duration = 5.seconds, promptHandler: PromptHandler = echoPromptHandler, + cancelHandler: CancelHandler = {}, block: suspend CoroutineScope.(TestAgent) -> Unit ) = runBlocking { val transport = TestTransport(timeout) val protocol = Protocol(this, transport) - val agentSupport = TestAgentSupport(promptHandler) + val agentSupport = TestAgentSupport(promptHandler, cancelHandler) val agent = Agent(protocol, agentSupport) protocol.start() val testAgent = TestAgent(agent, agentSupport, transport) @@ -187,10 +195,12 @@ fun withTestV2Agent( fun withInitializedTestAgent( timeout: Duration = 5.seconds, promptHandler: PromptHandler = echoPromptHandler, + cancelHandler: CancelHandler = {}, block: suspend CoroutineScope.(TestAgent) -> Unit ) = withTestAgent( timeout = timeout, promptHandler = promptHandler, + cancelHandler = cancelHandler, ) { testAgent -> testAgent.testInitialize(InitializeRequest(LATEST_PROTOCOL_VERSION)) check(testAgent.agentSupport.isInitialized) @@ -200,12 +210,14 @@ fun withInitializedTestAgent( fun withTestAgentSession( timeout: Duration = 5.seconds, promptHandler: PromptHandler = echoPromptHandler, + cancelHandler: CancelHandler = {}, cwd: String = ".", mcpServers: List = emptyList(), block: suspend CoroutineScope.(TestAgent, TestAgentSession) -> Unit ) = withInitializedTestAgent( timeout = timeout, promptHandler = promptHandler, + cancelHandler = cancelHandler, ) { testAgent -> val (newSessionResponse) = testAgent.testNewSession(NewSessionRequest(cwd, mcpServers)) checkNotNull(newSessionResponse) diff --git a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestTransport.kt b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestTransport.kt index 6aee87d..ba6a708 100644 --- a/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestTransport.kt +++ b/acp/src/jvmTest/kotlin/com/agentclientprotocol/agent/TestTransport.kt @@ -64,6 +64,14 @@ class TestTransport(val timeout: Duration) : BaseTransport() { } } + /** + * Fires a request without reading its reply, for tests that keep several requests in flight and read + * the replies with [receiveTestMessages]. + */ + fun fireTestRequestWithoutReply(methodName: MethodName, params: JsonElement, reqId: RequestId) { + fireMessage(JsonRpcRequest(reqId, methodName, params)) + } + fun fireTestNotification(methodName: MethodName, params: JsonElement) { fireMessage(JsonRpcNotification(methodName, params)) }