Skip to content

Commit c749a7a

Browse files
authored
Extensions API (#9)
* Implementation of extensions API. FileSystemOperations and TerminalOperations are implemented via extensions * Update README.md
1 parent 223e5af commit c749a7a

File tree

24 files changed

+1413
-500
lines changed

24 files changed

+1413
-500
lines changed

README.md

Lines changed: 266 additions & 174 deletions
Large diffs are not rendered by default.
Lines changed: 332 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,332 @@
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+
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,17 @@ import com.agentclientprotocol.agent.Agent
44
import com.agentclientprotocol.agent.AgentInfo
55
import com.agentclientprotocol.agent.AgentSession
66
import com.agentclientprotocol.agent.AgentSupport
7+
import com.agentclientprotocol.agent.client
78
import com.agentclientprotocol.client.Client
89
import com.agentclientprotocol.client.ClientInfo
910
import com.agentclientprotocol.client.ClientSession
1011
import com.agentclientprotocol.client.ClientSupport
1112
import com.agentclientprotocol.common.ClientSessionOperations
1213
import com.agentclientprotocol.common.Event
1314
import com.agentclientprotocol.common.SessionParameters
14-
import com.agentclientprotocol.common.client
1515
import com.agentclientprotocol.framework.ProtocolDriver
1616
import com.agentclientprotocol.model.ContentBlock
17+
import com.agentclientprotocol.model.LATEST_PROTOCOL_VERSION
1718
import com.agentclientprotocol.model.PermissionOption
1819
import com.agentclientprotocol.model.PermissionOptionId
1920
import com.agentclientprotocol.model.PermissionOptionKind
@@ -226,6 +227,7 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
226227
TODO("Not yet implemented")
227228
}
228229
})
230+
client.initialize(ClientInfo(protocolVersion = LATEST_PROTOCOL_VERSION))
229231
val session = client.newSession(SessionParameters("/test/path", emptyList()))
230232
val promptJob = launch {
231233
session.prompt(listOf(ContentBlock.Text("Test message"))).collect()
@@ -300,6 +302,7 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
300302
TODO("Not yet implemented")
301303
}
302304
})
305+
client.initialize(ClientInfo(protocolVersion = LATEST_PROTOCOL_VERSION))
303306
val session = client.newSession(SessionParameters("/test/path", emptyList()))
304307
val promptJob = launch {
305308
session.prompt(listOf(ContentBlock.Text("Test message"))).collect()
@@ -376,6 +379,7 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
376379
TODO("Not yet implemented")
377380
}
378381
})
382+
client.initialize(ClientInfo(protocolVersion = LATEST_PROTOCOL_VERSION))
379383
val responses = mutableListOf<String>()
380384
val session = client.newSession(SessionParameters("/test/path", emptyList()))
381385
session.prompt(listOf(ContentBlock.Text("Test message"))).collect { event ->
@@ -473,6 +477,7 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
473477
TODO("Not yet implemented")
474478
}
475479
})
480+
client.initialize(ClientInfo(protocolVersion = LATEST_PROTOCOL_VERSION))
476481
val responses = mutableListOf<String>()
477482
val session = client.newSession(SessionParameters("/test/path", emptyList()))
478483

@@ -571,6 +576,7 @@ abstract class SimpleAgentTest(protocolDriver: ProtocolDriver) : ProtocolDriver
571576
TODO("Not yet implemented")
572577
}
573578
})
579+
client.initialize(ClientInfo(protocolVersion = LATEST_PROTOCOL_VERSION))
574580
val responses = mutableListOf<String>()
575581
val session = client.newSession(SessionParameters("/test/path", emptyList()))
576582

0 commit comments

Comments
 (0)