Skip to content
Source
Sign inSign up

agent.test.ts

420 lines14.1 KB

1import * as Deferred from "effect/Deferred";
2import * as Effect from "effect/Effect";
3import * as Exit from "effect/Exit";
4import * as Fiber from "effect/Fiber";
5import * as Layer from "effect/Layer";
6import * as Queue from "effect/Queue";
7import * as Ref from "effect/Ref";
8import * as Schema from "effect/Schema";
9import * as Scope from "effect/Scope";
10​
11import { assert, it } from "@effect/vitest";
12​
13import * as AcpAgent from "./agent.ts";
14import * as AcpSchema from "./_generated/schema.gen.ts";
15import type * as AcpProtocol from "./protocol.ts";
16import {
17 encodeJsonl,
18 jsonRpcNotification,
19 jsonRpcRequest,
20 jsonRpcResponse,
21} from "./_internal/shared.ts";
22import { makeInMemoryStdio } from "./_internal/stdio.ts";
23​
24const RequestPermissionRequest = jsonRpcRequest(
25 "session/request_permission",
26 AcpSchema.RequestPermissionRequest,
27);
28const InitializeRequest = jsonRpcRequest("initialize", AcpSchema.InitializeRequest);
29const InitializeResponse = jsonRpcResponse(AcpSchema.InitializeResponse);
30const RequestPermissionResponse = jsonRpcResponse(AcpSchema.RequestPermissionResponse);
31const SessionCancelNotification = jsonRpcNotification(
32 "session/cancel",
33 AcpSchema.CancelSessionNotification,
34);
35const ExtPingNotification = jsonRpcNotification("x/ping", Schema.Struct({ count: Schema.Number }));
36const ExtRequest = jsonRpcRequest("x/test", Schema.Struct({ hello: Schema.String }));
37const ExtResponse = jsonRpcResponse(Schema.Struct({ ok: Schema.Boolean }));
38const DieResponse = Schema.Struct({
39 id: Schema.Number,
40 error: Schema.Struct({
41 code: Schema.Literal(-32603),
42 message: Schema.Literal("Internal error"),
43 }),
44});
45const decodeRequestPermissionRequest = Schema.decodeEffect(
46 Schema.fromJsonString(RequestPermissionRequest),
47);
48const decodeInitializeResponse = Schema.decodeEffect(Schema.fromJsonString(InitializeResponse));
49const decodeExtResponse = Schema.decodeEffect(Schema.fromJsonString(ExtResponse));
50​
51it.effect("effect-acp agent handles core agent requests and outbound client requests", () =>
52 Effect.gen(function* () {
53 const { stdio, input, output } = yield* makeInMemoryStdio();
54 const cancelNotifications = yield* Ref.make<Array<string>>([]);
55 const extNotifications = yield* Ref.make<Array<number>>([]);
56 const requestContexts = yield* Ref.make<Array<AcpProtocol.AcpRequestContext>>([]);
57 const cancelReceived = yield* Deferred.make<void>();
58 const extReceived = yield* Deferred.make<void>();
59 const scope = yield* Scope.make();
60 const context = yield* Layer.buildWithScope(AcpAgent.layer(stdio), scope);
61​
62 yield* Effect.gen(function* () {
63 const agent = yield* AcpAgent.AcpAgent;
64​
65 yield* agent.handleInitialize((_request, requestContext) =>
66 Ref.update(requestContexts, (current) => [...current, requestContext]).pipe(
67 Effect.as({
68 protocolVersion: 2,
69 capabilities: {},
70 info: {
71 name: "mock-agent",
72 version: "0.0.0",
73 },
74 }),
75 ),
76 );
77 yield* agent.handleExtRequest(
78 "x/test",
79 Schema.Struct({ hello: Schema.String }),
80 (_payload, requestContext) =>
81 Ref.update(requestContexts, (current) => [...current, requestContext]).pipe(
82 Effect.as({ ok: true }),
83 ),
84 );
85 yield* agent.handleCancel((notification) =>
86 Ref.update(cancelNotifications, (current) => [...current, notification.sessionId]).pipe(
87 Effect.andThen(Deferred.succeed(cancelReceived, undefined)),
88 ),
89 );
90 yield* agent.handleExtNotification(
91 "x/ping",
92 Schema.Struct({ count: Schema.Number }),
93 (payload) =>
94 Ref.update(extNotifications, (current) => [...current, payload.count]).pipe(
95 Effect.andThen(Deferred.succeed(extReceived, undefined)),
96 ),
97 );
98​
99 const permissionFiber = yield* agent.client
100 .requestPermission({
101 sessionId: "session-1",
102 title: "Allow mock action",
103 subject: {
104 type: "tool_call",
105 toolCall: {
106 toolCallId: "tool-1",
107 title: "Allow mock action",
108 },
109 },
110 options: [{ optionId: "allow", name: "Allow", kind: "allow_once" }],
111 })
112 .pipe(Effect.forkScoped);
113​
114 const permissionRequest = yield* decodeRequestPermissionRequest(yield* Queue.take(output));
115 assert.equal(permissionRequest.jsonrpc, "2.0");
116 assert.equal(permissionRequest.method, "session/request_permission");
117 assert.deepEqual(permissionRequest.params, {
118 sessionId: "session-1",
119 title: "Allow mock action",
120 subject: {
121 type: "tool_call",
122 toolCall: {
123 toolCallId: "tool-1",
124 title: "Allow mock action",
125 },
126 },
127 options: [{ optionId: "allow", name: "Allow", kind: "allow_once" }],
128 });
129 assert.deepEqual(permissionRequest.headers, []);
130​
131 yield* Queue.offer(
132 input,
133 yield* encodeJsonl(RequestPermissionResponse, {
134 jsonrpc: "2.0",
135 id: permissionRequest.id,
136 result: {
137 outcome: {
138 outcome: "selected",
139 optionId: "allow",
140 },
141 },
142 }),
143 );
144​
145 const permission = yield* Fiber.join(permissionFiber);
146 assert.equal(permission.outcome.outcome, "selected");
147​
148 yield* Queue.offer(
149 input,
150 yield* encodeJsonl(InitializeRequest, {
151 jsonrpc: "2.0",
152 id: 2,
153 method: "initialize",
154 params: {
155 protocolVersion: 2,
156 capabilities: {},
157 info: {
158 name: "effect-acp-test",
159 version: "0.0.0",
160 },
161 },
162 headers: [],
163 }),
164 );
165​
166 const initResponse = yield* decodeInitializeResponse(yield* Queue.take(output));
167 assert.deepEqual(initResponse, {
168 jsonrpc: "2.0",
169 id: 2,
170 result: {
171 protocolVersion: 2,
172 capabilities: {},
173 info: {
174 name: "mock-agent",
175 version: "0.0.0",
176 },
177 },
178 });
179​
180 yield* Queue.offer(
181 input,
182 yield* encodeJsonl(ExtRequest, {
183 jsonrpc: "2.0",
184 id: "extension-3",
185 method: "x/test",
186 params: { hello: "world" },
187 headers: [],
188 }),
189 );
190 const extResponse = yield* decodeExtResponse(yield* Queue.take(output));
191 assert.deepEqual(extResponse, {
192 jsonrpc: "2.0",
193 id: "extension-3",
194 result: { ok: true },
195 });
196 assert.deepEqual(yield* Ref.get(requestContexts), [
197 { requestId: "$t3:jsonrpc:number:2", method: "initialize" },
198 { requestId: "extension-3", method: "x/test" },
199 ]);
200​
201 yield* Queue.offer(
202 input,
203 yield* encodeJsonl(SessionCancelNotification, {
204 jsonrpc: "2.0",
205 method: "session/cancel",
206 params: {
207 sessionId: "session-1",
208 },
209 }),
210 );
211 yield* Queue.offer(
212 input,
213 yield* encodeJsonl(ExtPingNotification, {
214 jsonrpc: "2.0",
215 method: "x/ping",
216 params: { count: 2 },
217 }),
218 );
219​
220 yield* Deferred.await(cancelReceived);
221 yield* Deferred.await(extReceived);
222 assert.deepEqual(yield* Ref.get(cancelNotifications), ["session-1"]);
223 assert.deepEqual(yield* Ref.get(extNotifications), [2]);
224 }).pipe(Effect.provide(context), Effect.ensuring(Scope.close(scope, Exit.void)));
225 }),
226);
227​
228it.effect("effect-acp agent keeps RPC ids within int32 and apart from extension ids", () =>
229 Effect.gen(function* () {
230 const { stdio, input, output } = yield* makeInMemoryStdio();
231 const scope = yield* Scope.make();
232 const context = yield* Layer.buildWithScope(AcpAgent.layer(stdio), scope);
233​
234 yield* Effect.gen(function* () {
235 const agent = yield* AcpAgent.AcpAgent;
236​
237 const permissionFiber = yield* agent.client
238 .requestPermission({
239 sessionId: "session-1",
240 title: "Allow mock action",
241 subject: {
242 type: "tool_call",
243 toolCall: {
244 toolCallId: "tool-1",
245 title: "Allow mock action",
246 },
247 },
248 options: [{ optionId: "allow", name: "Allow", kind: "allow_once" }],
249 })
250 .pipe(Effect.forkScoped);
251 const extFiber = yield* agent.client
252 .extRequest("x/test", { hello: "world" })
253 .pipe(Effect.forkScoped);
254​
255 const firstOutbound = yield* Queue.take(output);
256 const secondOutbound = yield* Queue.take(output);
257​
258 const decodedPermission = Schema.decodeEffect(
259 Schema.fromJsonString(RequestPermissionRequest),
260 );
261 const decodedExt = Schema.decodeEffect(Schema.fromJsonString(ExtRequest));
262 const firstIsPermission = yield* decodedPermission(firstOutbound).pipe(Effect.isSuccess);
263​
264 const permissionRequest = firstIsPermission
265 ? yield* decodedPermission(firstOutbound)
266 : yield* decodedPermission(secondOutbound);
267 const extRequest = firstIsPermission
268 ? yield* decodedExt(secondOutbound)
269 : yield* decodedExt(firstOutbound);
270​
271 assert.notEqual(permissionRequest.id, extRequest.id);
272 // Keep locally generated ids compatible with SDKs that decode signed int32.
273 assert.typeOf(permissionRequest.id, "number");
274 assert.equal(permissionRequest.id, 2 ** 30);
275 assert.isAtMost(Number(permissionRequest.id), 2 ** 31 - 1);
276 assert.equal(extRequest.id, 1);
277​
278 yield* Queue.offer(
279 input,
280 yield* encodeJsonl(RequestPermissionResponse, {
281 jsonrpc: "2.0",
282 id: permissionRequest.id,
283 result: {
284 outcome: {
285 outcome: "selected",
286 optionId: "allow",
287 },
288 },
289 }),
290 );
291 yield* Queue.offer(
292 input,
293 yield* encodeJsonl(ExtResponse, {
294 jsonrpc: "2.0",
295 id: extRequest.id,
296 result: { ok: true },
297 }),
298 );
299​
300 const permission = yield* Fiber.join(permissionFiber);
301 assert.equal(permission.outcome.outcome, "selected");
302 assert.deepEqual(yield* Fiber.join(extFiber), { ok: true });
303​
304 const nextPermissionFiber = yield* agent.client
305 .requestPermission(permissionRequest.params)
306 .pipe(Effect.forkScoped);
307 const nextPermissionRequest = yield* decodedPermission(yield* Queue.take(output));
308 assert.equal(nextPermissionRequest.id, 2 ** 30 + 1);
309 yield* Queue.offer(
310 input,
311 yield* encodeJsonl(RequestPermissionResponse, {
312 jsonrpc: "2.0",
313 id: nextPermissionRequest.id,
314 result: { outcome: { outcome: "cancelled" } },
315 }),
316 );
317 assert.equal((yield* Fiber.join(nextPermissionFiber)).outcome.outcome, "cancelled");
318 }).pipe(Effect.provide(context), Effect.ensuring(Scope.close(scope, Exit.void)));
319 }),
320);
321​
322it.effect.each([0, 2 ** 30, 2 ** 32])(
323 "effect-acp agent answers incoming id %s independently of outbound ids",
324 (requestId) =>
325 Effect.gen(function* () {
326 const { stdio, input, output } = yield* makeInMemoryStdio();
327 const agent = yield* AcpAgent.make(stdio);
328 const contexts = yield* Ref.make<Array<AcpProtocol.AcpRequestContext>>([]);
329 yield* agent.handleInitialize((_request, context) =>
330 Ref.update(contexts, (current) => [...current, context]).pipe(
331 Effect.as({
332 protocolVersion: 2,
333 capabilities: {},
334 info: { name: "mock-agent", version: "0.0.0" },
335 }),
336 ),
337 );
338​
339 // Opposite directions may reuse an id even while a request is pending.
340 const outboundFiber = yield* agent.client
341 .requestPermission({
342 sessionId: "session-1",
343 title: "Allow mock action",
344 subject: {
345 type: "tool_call",
346 toolCall: { toolCallId: "tool-1", title: "Allow mock action" },
347 },
348 options: [{ optionId: "allow", name: "Allow", kind: "allow_once" }],
349 })
350 .pipe(Effect.forkScoped);
351 const outbound = yield* decodeRequestPermissionRequest(yield* Queue.take(output));
352​
353 yield* Queue.offer(
354 input,
355 yield* encodeJsonl(InitializeRequest, {
356 jsonrpc: "2.0",
357 id: requestId,
358 method: "initialize",
359 params: {
360 protocolVersion: 2,
361 capabilities: {},
362 info: { name: "test-client", version: "0.0.0" },
363 },
364 headers: [],
365 }),
366 );
367 assert.deepEqual(yield* decodeInitializeResponse(yield* Queue.take(output)), {
368 jsonrpc: "2.0",
369 id: requestId,
370 result: {
371 protocolVersion: 2,
372 capabilities: {},
373 info: { name: "mock-agent", version: "0.0.0" },
374 },
375 });
376 assert.deepEqual(yield* Ref.get(contexts), [
377 { requestId: `$t3:jsonrpc:number:${requestId}`, method: "initialize" },
378 ]);
379​
380 yield* Queue.offer(
381 input,
382 yield* encodeJsonl(RequestPermissionResponse, {
383 jsonrpc: "2.0",
384 id: outbound.id,
385 result: { outcome: { outcome: "cancelled" } },
386 }),
387 );
388 assert.equal((yield* Fiber.join(outboundFiber)).outcome.outcome, "cancelled");
389 }),
390);
391​
392it.effect("effect-acp agent answers a request whose handler dies with an error for it", () =>
393 Effect.gen(function* () {
394 const { stdio, input, output } = yield* makeInMemoryStdio();
395 const scope = yield* Scope.make();
396 const context = yield* Layer.buildWithScope(AcpAgent.layer(stdio), scope);
397 const agent = yield* Effect.service(AcpAgent.AcpAgent).pipe(Effect.provide(context));
398 yield* agent.handleInitialize(() => Effect.die(new Error("handler bug")));
399​
400 yield* Queue.offer(
401 input,
402 yield* encodeJsonl(InitializeRequest, {
403 jsonrpc: "2.0",
404 id: 7,
405 method: "initialize",
406 params: {
407 protocolVersion: 2,
408 capabilities: {},
409 info: { name: "effect-acp-test", version: "0.0.0" },
410 },
411 headers: [],
412 }),
413 );
414 const response = yield* Queue.take(output).pipe(
415 Effect.flatMap(Schema.decodeEffect(Schema.fromJsonString(DieResponse))),
416 );
417 assert.equal(response.id, 7);
418 yield* Scope.close(scope, Exit.void);
419 }),
420);