refactor: migrate to @bufbuild/protobuf for JS protobufs (#295)

* refactor: migrate to @bufbuild/protobuf for JS protobufs

* bump versions

Author: grant0417

packages/api-bindings-wrappers/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws/amazon-q-developer-cli-api-bindings-wrappers",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "license": "MIT OR Apache-2.0",
   "author": "Amazon Web Services",
   "repository": "https://github.com/aws/amazon-q-developer-cli",

packages/api-bindings/codegen/generate-requests.ts

@@ -1,10 +1,5 @@
-import {
-  Project,
-  PropertySignature,
-  SourceFile,
-  CodeBlockWriter,
-  IndentationText,
-} from "ts-morph";
+import { file_fig as file } from "@aws/amazon-q-developer-cli-proto/fig";
+import { CodeBlockWriter, IndentationText, Project } from "ts-morph";
 import { readFileSync } from "node:fs";
 import { join, dirname } from "node:path";
 import { fileURLToPath } from "node:url";
@@ -30,21 +25,13 @@ const normalize = (type: string): string => {
   return capitalizeFirstLetter(normalized);
 };
 
-const getSubmessageTypes = (bindings: SourceFile, interfaceName: string) => {
-  const interfaceRef = bindings.getInterface(interfaceName)!;
-  const submessage = interfaceRef.getProperties()[1];
-
-  const submessageUnion = submessage
-    .getChildren()
-    .filter((elm) => elm.getKindName() === "UnionType")[0];
-  const literals = submessageUnion
-    .getChildren()[0]
-    .getChildren()
-    .filter((elm) => elm.getKindName() === "TypeLiteral");
-  const types = literals
-    .map((elm) => elm.getChildren()[1])
-    .map((elm) => elm.getChildren()[1]);
-  return types.map((prop) => (prop as PropertySignature).getName());
+const getSubmessageTypes = (interfaceName: string) => {
+  const message = file.messages.find(
+    (message) => message.name === interfaceName,
+  );
+  return (
+    message?.fields.map((type) => type.message?.name!).filter(Boolean) ?? []
+  );
 };
 
 const writeGenericSendRequestWithResponseFunction = (
@@ -54,23 +41,23 @@ const writeGenericSendRequestWithResponseFunction = (
   const lowercasedEndpoint = lowercaseFirstLetter(endpoint);
 
   const template = `export async function send${endpoint}Request(
-request: ${endpoint}Request
+request: Omit<${endpoint}Request, "$typeName" | "$unknown">
 ): Promise<${endpoint}Response> {
   return new Promise((resolve, reject) => {
     sendMessage(
-      { $case: "${lowercasedEndpoint}Request", ${lowercasedEndpoint}Request: request },
+      { case: "${lowercasedEndpoint}Request", value: create(${endpoint}RequestSchema, request) },
       (response) => {
-        switch (response?.$case) {
+        switch (response?.case) {
           case "${lowercasedEndpoint}Response":
-            resolve(response.${lowercasedEndpoint}Response);
+            resolve(response.value);
             break;
           case "error":
-            reject(Error(response.error));
+            reject(Error(response.value));
             break;
           default:
             reject(
               Error(
-                  \`Invalid response '\${response?.$case}' for '${endpoint}Request'\`
+                  \`Invalid response '\${response?.case}' for '${endpoint}Request'\`
               )
             );
         }
@@ -89,23 +76,23 @@ const writeGenericSendRequestFunction = (
   const lowercasedEndpoint = lowercaseFirstLetter(endpoint);
 
   const template = `export async function send${endpoint}Request(
-  request: ${endpoint}Request
+  request: Omit<${endpoint}Request, "$typeName" | "$unknown">
 ): Promise<void> {
   return new Promise((resolve, reject) => {
     sendMessage(
-      { $case: "${lowercasedEndpoint}Request", ${lowercasedEndpoint}Request: request },
+      { case: "${lowercasedEndpoint}Request", value: create(${endpoint}RequestSchema, request) },
       (response) => {
-        switch (response?.$case) {
+        switch (response?.case) {
           case "success":
             resolve();
             break;
           case "error":
-            reject(Error(response.error));
+            reject(Error(response.value));
             break;
           default:
             reject(
               Error(
-                \`Invalid response '\${response?.$case}' for '${endpoint}Request'\`
+                \`Invalid response '\${response?.case}' for '${endpoint}Request'\`
               )
             );
         }
@@ -124,28 +111,20 @@ const project = new Project({
 
 project.addSourceFilesAtPaths(join(__dirname, "../src/*.ts"));
 
-const text = readFileSync(
-  "node_modules/@aws/amazon-q-developer-cli-proto/dist/fig.pb.ts",
-  "utf8",
-);
-const protobufBindings = project.createSourceFile("fig.pb.ts", text);
-
-const requestTypes = getSubmessageTypes(
-  protobufBindings,
-  "ClientOriginatedMessage",
+const requestTypes = getSubmessageTypes("ClientOriginatedMessage");
+const responseTypes = getSubmessageTypes("ServerOriginatedMessage").filter(
+  (type) => type.includes("Response"),
 );
-const responseTypes = getSubmessageTypes(
-  protobufBindings,
-  "ServerOriginatedMessage",
-).filter((type) => type.includes("Response"));
 
 const [requestsWithMatchingResponses, otherRequests] = requestTypes
   .filter((request) => request !== "notificationRequest")
   .reduce(
     (result, request) => {
       const [matchingResponse, other] = result;
 
-      const endpoint = lowercaseFirstLetter(normalize(request));
+      const endpoint = normalize(request);
+
+      console.log(endpoint, requestTypes);
 
       if (responseTypes.indexOf(`${endpoint}Response`) !== -1) {
         return [matchingResponse.concat([request]), other];
@@ -175,15 +154,31 @@ const sourceFile = project.createSourceFile(
     const responses = requestsWithMatchingResponses.map((request) =>
       request.replace("Request", "Response"),
     );
-    const imports = requestsWithMatchingResponses
-      .concat(responses)
-      .concat(otherRequests)
+    const requestSchemas = requestsWithMatchingResponses.map(
+      (s) => `${s}Schema`,
+    );
+    const otherRequestSchemas = otherRequests.map((s) => `${s}Schema`);
+    const imports = [
+      requestsWithMatchingResponses,
+      responses,
+      otherRequests,
+      requestSchemas,
+      otherRequestSchemas,
+    ]
+      .flat()
       .sort()
       .map(capitalizeFirstLetter);
     writer.writeLine(
       `import { \n${imports.join(",\n")}\n } from "@aws/amazon-q-developer-cli-proto/fig";`,
     );
-    writer.writeLine(`import { sendMessage } from "./core.js";`).blankLine();
+    writer
+      .writeLine(
+        [
+          'import { sendMessage } from "./core.js";',
+          'import { create } from "@bufbuild/protobuf";',
+        ].join("\n"),
+      )
+      .blankLine();
 
     requestsWithMatchingResponses.forEach((request) =>
       writeGenericSendRequestWithResponseFunction(writer, normalize(request)),

packages/api-bindings/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws/amazon-q-developer-cli-api-bindings",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "license": "MIT OR Apache-2.0",
   "author": "Amazon Web Services",
   "repository": "https://github.com/aws/amazon-q-developer-cli",
@@ -22,7 +22,8 @@
     "generate-requests": "tsx codegen/generate-requests.ts && prettier -w src/requests.ts"
   },
   "dependencies": {
-    "@aws/amazon-q-developer-cli-proto": "workspace:^"
+    "@aws/amazon-q-developer-cli-proto": "workspace:^",
+    "@bufbuild/protobuf": "2.2.3"
   },
   "devDependencies": {
     "@amzn/eslint-config": "workspace:^",
@@ -34,8 +35,6 @@
     "lint-staged": "^15.2.10",
     "prettier": "^3.4.2",
     "ts-morph": "^24.0.0",
-    "ts-proto": "~2.5.0",
-    "tslib": "^2.8.1",
     "tsx": "^4.19.2",
     "typescript": "^5.7.2"
   },

packages/api-bindings/src/auth.ts

@@ -54,7 +54,10 @@ export function startPkceAuthorization({
 
 export function finishPkceAuthorization({
   authRequestId,
-}: AuthFinishPkceAuthorizationRequest): Promise<AuthFinishPkceAuthorizationResponse> {
+}: Omit<
+  AuthFinishPkceAuthorizationRequest,
+  "$typeName"
+>): Promise<AuthFinishPkceAuthorizationResponse> {
   return sendAuthFinishPkceAuthorizationRequest({ authRequestId });
 }
 

packages/api-bindings/src/core.ts

@@ -1,9 +1,12 @@
 import {
-  ServerOriginatedMessage,
-  ClientOriginatedMessage,
+  type ServerOriginatedMessage,
+  type ClientOriginatedMessage,
+  ClientOriginatedMessageSchema,
+  ServerOriginatedMessageSchema,
 } from "@aws/amazon-q-developer-cli-proto/fig";
 
 import { b64ToBytes, bytesToBase64 } from "./utils.js";
+import { create, fromBinary, toBinary } from "@bufbuild/protobuf";
 
 interface GlobalAPIError {
   error: string;
@@ -19,9 +22,9 @@ export type APIResponseHandler = (
 ) => shouldKeepListening | void;
 
 let messageId = 0;
-const handlers: Record<number, APIResponseHandler> = {};
+const handlers: Record<string, APIResponseHandler> = {};
 
-export function setHandlerForId(handler: APIResponseHandler, id: number) {
+export function setHandlerForId(handler: APIResponseHandler, id: string) {
   handlers[id] = handler;
 }
 
@@ -30,45 +33,46 @@ const receivedMessage = (response: ServerOriginatedMessage): void => {
     return;
   }
 
-  const handler = handlers[response.id];
+  const id = response.id.toString();
+  const handler = handlers[id];
 
   if (!handler) {
     return;
   }
 
-  const keepListeningOnID = handlers[response.id](response.submessage);
+  const keepListeningOnID = handlers[id](response.submessage);
 
   if (!keepListeningOnID) {
-    delete handlers[response.id];
+    delete handlers[id];
   }
 };
 
 export function sendMessage(
-  message: ClientOriginatedMessage["submessage"],
+  submessage: ClientOriginatedMessage["submessage"],
   handler?: APIResponseHandler,
 ) {
-  const request: ClientOriginatedMessage = {
-    id: (messageId += 1),
-    submessage: message,
-  };
+  const request: ClientOriginatedMessage = create(
+    ClientOriginatedMessageSchema,
+    {
+      id: BigInt((messageId += 1)),
+      submessage,
+    },
+  );
 
   if (handler && request.id) {
-    handlers[request.id] = handler;
+    handlers[request.id.toString()] = handler;
   }
 
-  const buffer = ClientOriginatedMessage.encode(request).finish();
+  const buffer = toBinary(ClientOriginatedMessageSchema, request);
 
   if (
     window?.fig?.constants?.supportApiProto &&
     window?.fig?.constants?.apiProtoUrl
   ) {
     const url = new URL(window.fig.constants.apiProtoUrl);
 
-    if (
-      request.submessage?.$case &&
-      typeof request.submessage?.$case === "string"
-    ) {
-      url.pathname = `/${request.submessage.$case}`;
+    if (typeof request.submessage?.case === "string") {
+      url.pathname = `/${request.submessage.case}`;
     } else {
       url.pathname = "/unknown";
     }
@@ -81,8 +85,8 @@ export function sendMessage(
       body: buffer,
     }).then(async (res) => {
       const body = new Uint8Array(await res.arrayBuffer());
-      const m = ServerOriginatedMessage.decode(body);
-      receivedMessage(m);
+      const message = fromBinary(ServerOriginatedMessageSchema, body);
+      receivedMessage(message);
     });
     return;
   }
@@ -115,7 +119,7 @@ const setupEventListeners = (): void => {
   document.addEventListener(FigProtoMessageReceived, (event: Event) => {
     const raw = (event as CustomEvent).detail as string;
     const bytes = b64ToBytes(raw);
-    const message = ServerOriginatedMessage.decode(bytes);
+    const message = fromBinary(ServerOriginatedMessageSchema, bytes);
     receivedMessage(message);
   });
 };

packages/api-bindings/src/editbuffer.ts

@@ -12,9 +12,9 @@ export function subscribe(
   return _subscribe(
     { type: NotificationType.NOTIFY_ON_EDITBUFFFER_CHANGE },
     (notification) => {
-      switch (notification?.type?.$case) {
+      switch (notification?.type?.case) {
         case "editBufferNotification":
-          return handler(notification.type.editBufferNotification);
+          return handler(notification.type.value);
         default:
           break;
       }

packages/api-bindings/src/event.ts

@@ -8,10 +8,9 @@ export function subscribe<T>(
   return _subscribe(
     { type: NotificationType.NOTIFY_ON_EVENT },
     (notification) => {
-      switch (notification?.type?.$case) {
+      switch (notification?.type?.case) {
         case "eventNotification": {
-          const { eventName: name, payload } =
-            notification.type.eventNotification;
+          const { eventName: name, payload } = notification.type.value;
           if (name === eventName) {
             try {
               return handler(payload ? JSON.parse(payload) : null);