1+ package com.agentclientprotocol
2+
3+ import com.agentclientprotocol.agent.Agent
4+ import com.agentclientprotocol.agent.AgentInfo
5+ import com.agentclientprotocol.agent.AgentSession
6+ import com.agentclientprotocol.agent.AgentSupport
7+ import com.agentclientprotocol.client.*
8+ import com.agentclientprotocol.common.*
9+ import com.agentclientprotocol.framework.ProtocolDriver
10+ import com.agentclientprotocol.model.*
11+ import com.agentclientprotocol.protocol.RpcMethodsOperations
12+ import com.agentclientprotocol.protocol.invoke
13+ import kotlinx.coroutines.currentCoroutineContext
14+ import kotlinx.coroutines.flow.Flow
15+ import kotlinx.coroutines.flow.flow
16+ import kotlinx.coroutines.test.TestResult
17+ import kotlinx.serialization.json.JsonElement
18+ import kotlin.test.*
19+
20+ open class TestClientSessionOperations (): ClientSessionOperations {
21+ override suspend fun requestPermissions (
22+ toolCall : SessionUpdate .ToolCallUpdate ,
23+ permissions : List <PermissionOption >,
24+ _meta : JsonElement ? ,
25+ ): RequestPermissionResponse {
26+ TODO ()
27+ }
28+
29+ override suspend fun notify (
30+ notification : SessionUpdate ,
31+ _meta : JsonElement ? ,
32+ ) {
33+ TODO ()
34+ }
35+ }
36+
37+ class TestClientSupport (val createSessionFunc : suspend (ClientSession , JsonElement ? ) -> ClientSessionOperations ) : ClientSupport {
38+ override suspend fun createClientSession (
39+ session : ClientSession ,
40+ _sessionResponseMeta : JsonElement ? ,
41+ ): ClientSessionOperations {
42+ return createSessionFunc(session, _sessionResponseMeta )
43+ }
44+ }
45+
46+ open class TestAgentSession (override val sessionId : SessionId = SessionId ("test-session-id")) : AgentSession {
47+ override suspend fun prompt (
48+ content : List <ContentBlock >,
49+ _meta : JsonElement ? ,
50+ ): Flow <Event > {
51+ TODO ()
52+ }
53+ }
54+
55+ class TestAgentSupport (val capabilities : AgentCapabilities = AgentCapabilities (), val createSessionFunc : suspend (SessionParameters ) -> AgentSession ) : AgentSupport {
56+ val agentInitialized = kotlinx.coroutines.CompletableDeferred <ClientInfo >()
57+
58+ override suspend fun initialize (clientInfo : ClientInfo ): AgentInfo {
59+ agentInitialized.complete(clientInfo)
60+ return AgentInfo (clientInfo.protocolVersion, capabilities)
61+ }
62+
63+ override suspend fun createSession (sessionParameters : SessionParameters ): AgentSession {
64+ return createSessionFunc(sessionParameters)
65+ }
66+
67+ override suspend fun loadSession (
68+ sessionId : SessionId ,
69+ sessionParameters : SessionParameters ,
70+ ): AgentSession {
71+ return createSessionFunc(sessionParameters)
72+ }
73+ }
74+
75+
76+ abstract class ExtensionsTest (protocolDriver : ProtocolDriver ) : ProtocolDriver by protocolDriver {
77+
78+ @Test
79+ fun `exception when extension interface not supported on client` () = testWithProtocols { clientProtocol, agentProtocol ->
80+
81+ val clientSupport = TestClientSupport { session, _sessionResponseMeta ->
82+ return @TestClientSupport TestClientSessionOperations ()
83+ }
84+ val client = Client (protocol = clientProtocol, clientSupport = clientSupport, handlerSideExtensions = listOf (
85+ FileSystemOperations ))
86+
87+ val agentSupport = TestAgentSupport {
88+ object : TestAgentSession () {
89+ override suspend fun prompt (content : List <ContentBlock >, _meta : JsonElement ? ): Flow <Event > = flow {
90+ emit(Event .SessionUpdateEvent (agentTextChunk(" Hello" )))
91+ val fileSystemOperations =
92+ currentCoroutineContext().remoteSessionOperations(FileSystemOperations )
93+ val readTextFileResponse = fileSystemOperations.fsReadTextFile(" /test/path" )
94+ emit(Event .SessionUpdateEvent (agentTextChunk(" File content: ${readTextFileResponse.content} " )))
95+ }
96+ }
97+ }
98+ val agent = Agent (agentProtocol, agentSupport, remoteSideExtensions = listOf (FileSystemOperations ))
99+
100+ client.initialize(
101+ ClientInfo (
102+ capabilities = ClientCapabilities (
103+ fs = FileSystemCapability (
104+ readTextFile = true ,
105+ writeTextFile = true
106+ )
107+ )
108+ )
109+ )
110+ val session = client.newSession(SessionParameters (" /test/path" , emptyList()))
111+ val exception = runCatching { session.prompt(textBlocks(" Test" )).collect { } }.exceptionOrNull()
112+ assertNotNull(exception, " Exception should be thrown" )
113+ assertContains(exception.message!! , " does not implement extension type" )
114+ }
115+
116+ @Test
117+ fun `exception when extension not registered on client` () = testWithProtocols { clientProtocol, agentProtocol ->
118+
119+ val clientSupport = TestClientSupport { session, _sessionResponseMeta ->
120+ return @TestClientSupport TestClientSessionOperations ()
121+ }
122+ val client = Client (protocol = clientProtocol, clientSupport = clientSupport)
123+
124+ val agentSupport = TestAgentSupport {
125+ object : TestAgentSession () {
126+ override suspend fun prompt (content : List <ContentBlock >, _meta : JsonElement ? ): Flow <Event > = flow {
127+ emit(Event .SessionUpdateEvent (agentTextChunk(" Hello" )))
128+ val fileSystemOperations =
129+ currentCoroutineContext().remoteSessionOperations(FileSystemOperations )
130+ val readTextFileResponse = fileSystemOperations.fsReadTextFile(" /test/path" )
131+ emit(Event .SessionUpdateEvent (agentTextChunk(" File content: ${readTextFileResponse.content} " )))
132+ }
133+ }
134+ }
135+ val agent = Agent (agentProtocol, agentSupport, remoteSideExtensions = listOf (FileSystemOperations ))
136+
137+ client.initialize(
138+ ClientInfo (
139+ capabilities = ClientCapabilities (
140+ fs = FileSystemCapability (
141+ readTextFile = true ,
142+ writeTextFile = true
143+ )
144+ )
145+ )
146+ )
147+ val session = client.newSession(SessionParameters (" /test/path" , emptyList()))
148+ val exception = runCatching { session.prompt(textBlocks(" Test" )).collect { } }.exceptionOrNull()
149+ assertNotNull(exception, " Exception should be thrown" )
150+ assertContains(exception.message!! , " Method not supported" )
151+ }
152+
153+ @Test
154+ fun `exception when extension not registered on agent` () = testWithProtocols { clientProtocol, agentProtocol ->
155+
156+ val clientSupport = TestClientSupport { session, _sessionResponseMeta ->
157+ return @TestClientSupport object : TestClientSessionOperations (), FileSystemOperations {
158+ override suspend fun fsReadTextFile (
159+ path : String ,
160+ line : UInt? ,
161+ limit : UInt? ,
162+ _meta : JsonElement ? ,
163+ ): ReadTextFileResponse {
164+ return ReadTextFileResponse (" test" )
165+ }
166+
167+ override suspend fun fsWriteTextFile (
168+ path : String ,
169+ content : String ,
170+ _meta : JsonElement ? ,
171+ ): WriteTextFileResponse {
172+ return WriteTextFileResponse ()
173+ }
174+ }
175+ }
176+ val client = Client (protocol = clientProtocol, clientSupport = clientSupport, handlerSideExtensions = listOf (
177+ FileSystemOperations )
178+ )
179+
180+ val agentSupport = TestAgentSupport {
181+ object : TestAgentSession () {
182+ override suspend fun prompt (content : List <ContentBlock >, _meta : JsonElement ? ): Flow <Event > = flow {
183+ emit(Event .SessionUpdateEvent (agentTextChunk(" Hello" )))
184+ val fileSystemOperations =
185+ currentCoroutineContext().remoteSessionOperations(FileSystemOperations )
186+ val readTextFileResponse = fileSystemOperations.fsReadTextFile(" /test/path" )
187+ emit(Event .SessionUpdateEvent (agentTextChunk(" File content: ${readTextFileResponse.content} " )))
188+ }
189+ }
190+ }
191+ val agent = Agent (agentProtocol, agentSupport, /* remoteSideExtensions = listOf(FileSystemOperations)*/ )
192+
193+ client.initialize(
194+ ClientInfo (
195+ capabilities = ClientCapabilities (
196+ fs = FileSystemCapability (
197+ readTextFile = true ,
198+ writeTextFile = true
199+ )
200+ )
201+ )
202+ )
203+ val session = client.newSession(SessionParameters (" /test/path" , emptyList()))
204+ val exception = runCatching { session.prompt(textBlocks(" Test" )).collect { } }.exceptionOrNull()
205+ assertNotNull(exception, " Exception should be thrown" )
206+ assertContains(exception.message!! , " is either not registered" )
207+ }
208+
209+ @Test
210+ fun `call client extension from agent` () = testWithProtocols { clientProtocol, agentProtocol ->
211+
212+ val clientSupport = TestClientSupport { session, _sessionResponseMeta ->
213+ return @TestClientSupport object : TestClientSessionOperations (), FileSystemOperations {
214+ override suspend fun fsReadTextFile (
215+ path : String ,
216+ line : UInt? ,
217+ limit : UInt? ,
218+ _meta : JsonElement ? ,
219+ ): ReadTextFileResponse {
220+ return ReadTextFileResponse (" test" )
221+ }
222+
223+ override suspend fun fsWriteTextFile (
224+ path : String ,
225+ content : String ,
226+ _meta : JsonElement ? ,
227+ ): WriteTextFileResponse {
228+ return WriteTextFileResponse ()
229+ }
230+ }
231+ }
232+ val client = Client (protocol = clientProtocol, clientSupport = clientSupport, handlerSideExtensions = listOf (
233+ FileSystemOperations )
234+ )
235+
236+ val agentSupport = TestAgentSupport {
237+ object : TestAgentSession () {
238+ override suspend fun prompt (content : List <ContentBlock >, _meta : JsonElement ? ): Flow <Event > = flow {
239+ emit(Event .SessionUpdateEvent (agentTextChunk(" Hello" )))
240+ val fileSystemOperations =
241+ currentCoroutineContext().remoteSessionOperations(FileSystemOperations )
242+ val readTextFileResponse = fileSystemOperations.fsReadTextFile(" /test/path" )
243+ emit(Event .SessionUpdateEvent (agentTextChunk(" File content: ${readTextFileResponse.content} " )))
244+ }
245+ }
246+ }
247+ val agent = Agent (agentProtocol, agentSupport, remoteSideExtensions = listOf (FileSystemOperations ))
248+
249+ client.initialize(
250+ ClientInfo (
251+ capabilities = ClientCapabilities (
252+ fs = FileSystemCapability (
253+ readTextFile = true ,
254+ writeTextFile = true
255+ )
256+ )
257+ )
258+ )
259+ val session = client.newSession(SessionParameters (" /test/path" , emptyList()))
260+ val result = session.promptToList(" Test" )
261+ assertContentEquals(listOf (" Hello" , " File content: test" , " END_TURN" ), result)
262+ }
263+
264+ interface TestAgentInterface {
265+ suspend fun readTextFile (path : String ): String
266+
267+ companion object : HandlerSideExtension <TestAgentInterface >, RemoteSideExtension <TestAgentInterface > {
268+ override fun RegistrarContext<TestAgentInterface>.register () {
269+ setSessionExtensionRequestHandler(AcpMethod .ClientMethods .FsReadTextFile ) { ops, params ->
270+ return @setSessionExtensionRequestHandler ReadTextFileResponse (ops.readTextFile(params.path))
271+ }
272+ }
273+
274+ override val name: String
275+ get() = TestAgentInterface ::class .simpleName!!
276+
277+ override fun isSupported (remoteSideCapabilities : AcpCapabilities ): Boolean {
278+ return true
279+ }
280+
281+ override fun createSessionRemote (
282+ rpc : RpcMethodsOperations ,
283+ capabilities : AcpCapabilities ,
284+ sessionId : SessionId ,
285+ ): TestAgentInterface {
286+ return object : TestAgentSession (), TestAgentInterface {
287+ override suspend fun readTextFile (path : String ): String = AcpMethod .ClientMethods .FsReadTextFile (rpc,
288+ ReadTextFileRequest (sessionId, path)).content
289+ }
290+ }
291+ }
292+ }
293+
294+ @Test
295+ fun `call agent extension from client` (): TestResult = testWithProtocols { clientProtocol, agentProtocol ->
296+
297+ val clientSupport = TestClientSupport { session, _sessionResponseMeta ->
298+ return @TestClientSupport object : TestClientSessionOperations () {
299+ }
300+ }
301+ val client = Client (
302+ protocol = clientProtocol, clientSupport = clientSupport, remoteSideExtensions = listOf (
303+ TestAgentInterface
304+ )
305+ )
306+
307+ val agentSupport = TestAgentSupport {
308+ object : TestAgentSession (), TestAgentInterface {
309+ override suspend fun readTextFile (path : String ): String = " test content"
310+ }
311+ }
312+ val agent = Agent (agentProtocol, agentSupport, handlerSideExtensions = listOf (TestAgentInterface ))
313+
314+ client.initialize(
315+ ClientInfo (
316+ capabilities = ClientCapabilities (
317+ fs = FileSystemCapability (
318+ readTextFile = true ,
319+ writeTextFile = true
320+ )
321+ )
322+ )
323+ )
324+ val session = client.newSession(SessionParameters (" /test/path" , emptyList()))
325+ val fileSystemOperations = session.remoteOperations(TestAgentInterface )
326+ val fileResponse = fileSystemOperations.readTextFile(" /test/file/path" )
327+ assertEquals(" test content" , fileResponse)
328+ }
329+
330+
331+
332+ }
0 commit comments