Skip to content

Commit

Permalink
Make it easy to issue progress notifications from request handlers
Browse files Browse the repository at this point in the history
Resolves #34.
  • Loading branch information
jspahrsummers committed Nov 4, 2024
1 parent 822c451 commit 231026e
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 5 deletions.
33 changes: 29 additions & 4 deletions src/shared/protocol.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { AnyZodObject, ZodLiteral, ZodObject, z } from "zod";
import { AnyZodObject, ZodLiteral, ZodObject, ZodOptional, z } from "zod";
import {
BaseRequestParamsSchema,
ErrorCode,
JSONRPCError,
JSONRPCNotification,
Expand Down Expand Up @@ -312,18 +313,42 @@ export class Protocol<
/**
* Registers a handler to invoke when this protocol object receives a request with the given method.
*
* The handler receives a second callback parameter that can be used to emit progress notifications.
*
* Note that this will replace any previous request handler for the same method.
*/
setRequestHandler<
T extends ZodObject<{
method: ZodLiteral<string>;
params: ZodOptional<typeof BaseRequestParamsSchema>;
}>,
>(
requestSchema: T,
handler: (request: z.infer<T>) => SendResultT | Promise<SendResultT>,
handler: (
request: z.infer<T>,
progress: (progress: Progress) => void,
) => SendResultT | Promise<SendResultT>,
): void {
this._requestHandlers.set(requestSchema.shape.method.value, (request) =>
Promise.resolve(handler(requestSchema.parse(request))),
this._requestHandlers.set(
requestSchema.shape.method.value,
async (request) => {
const parsedRequest = requestSchema.parse(request);
const progressToken = parsedRequest.params?._meta?.progressToken;
const progressHandler =
progressToken !== undefined
? (progress: Progress) =>
// Sending directly on the transport to avoid typing conflicts
this._transport
?.send({
jsonrpc: "2.0",
method: "notifications/progress",
params: { ...progress, progressToken },
})
.catch((error: Error) => this._onerror(error))
: () => {};

return handler(parsedRequest, progressHandler);
},
);
}

Expand Down
2 changes: 1 addition & 1 deletion src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ export const ProgressTokenSchema = z.union([z.string(), z.number().int()]);
*/
export const CursorSchema = z.string();

const BaseRequestParamsSchema = z
export const BaseRequestParamsSchema = z
.object({
_meta: z.optional(
z
Expand Down

0 comments on commit 231026e

Please sign in to comment.