Feat: Reuse signed request when reading state

packages/agent/src/agent/api.ts

@@ -111,16 +111,32 @@ export interface Agent {
    */
   getPrincipal(): Promise<Principal>;
 
+  /**
+   * Create the request for the read state call.
+   * `readState` uses this internally.
+   * Useful to avoid signing the same request multiple times.
+   */
+  createReadStateRequest?(
+    options: ReadStateOptions,
+    identity?: Identity,
+    // eslint-disable-next-line  @typescript-eslint/no-explicit-any
+  ): Promise<any>;
+
   /**
    * Send a read state query to the replica. This includes a list of paths to return,
    * and will return a Certificate. This will only reject on communication errors,
    * but the certificate might contain less information than requested.
    * @param effectiveCanisterId A Canister ID related to this call.
    * @param options The options for this call.
+   * @param identity Identity for the call. If not specified, uses the instance identity.
+   * @param request The request to send in case it has already been created.
    */
   readState(
     effectiveCanisterId: Principal | string,
     options: ReadStateOptions,
+    identity?: Identity,
+    // eslint-disable-next-line  @typescript-eslint/no-explicit-any
+    request?: any,
   ): Promise<ReadStateResponse>;
 
   call(canisterId: Principal | string, fields: CallOptions): Promise<SubmitResponse>;

packages/agent/src/agent/http/http.test.ts

@@ -1,7 +1,13 @@
 import { HttpAgent, Nonce } from '../index';
 import * as cbor from '../../cbor';
 import { Expiry, makeNonceTransform } from './transforms';
-import { CallRequest, Envelope, makeNonce, SubmitRequestType } from './types';
+import {
+  CallRequest,
+  Envelope,
+  HttpAgentRequestTransformFn,
+  makeNonce,
+  SubmitRequestType,
+} from './types';
 import { Principal } from '@dfinity/principal';
 import { requestIdOf } from '../../request_id';
 
@@ -154,6 +160,59 @@ test('queries with the same content should have the same signature', async () =>
   expect(response3).toEqual(response4);
 });
 
+test('readState should not call transformers if request is passed', async () => {
+  const mockResponse = {
+    status: 'replied',
+    reply: { arg: new Uint8Array([]) },
+  };
+
+  const mockFetch: jest.Mock = jest.fn((resource, init) => {
+    const body = cbor.encode(mockResponse);
+    return Promise.resolve(
+      new Response(body, {
+        status: 200,
+      }),
+    );
+  });
+
+  const canisterIdent = '2chl6-4hpzw-vqaaa-aaaaa-c';
+  const nonce = new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7]) as Nonce;
+
+  const principal = await Principal.anonymous();
+
+  const httpAgent = new HttpAgent({
+    fetch: mockFetch,
+    host: 'http://localhost',
+    disableNonce: true,
+  });
+  httpAgent.addTransform(makeNonceTransform(() => nonce));
+  const transformMock: HttpAgentRequestTransformFn = jest
+    .fn()
+    .mockImplementation(d => Promise.resolve(d));
+  httpAgent.addTransform(transformMock);
+
+  const methodName = 'greet';
+  const arg = new Uint8Array([]);
+
+  const requestId = await requestIdOf({
+    request_type: SubmitRequestType.Call,
+    nonce,
+    canister_id: Principal.fromText(canisterIdent).toString(),
+    method_name: methodName,
+    arg,
+    sender: principal,
+  });
+
+  const paths = [
+    [new TextEncoder().encode('request_status'), requestId, new TextEncoder().encode('reply')],
+  ];
+
+  const request = await httpAgent.createReadStateRequest({ paths });
+  expect(transformMock).toBeCalledTimes(1);
+  await httpAgent.readState(canisterIdent, { paths }, undefined, request);
+  expect(transformMock).toBeCalledTimes(1);
+});
+
 test('redirect avoid', async () => {
   function checkUrl(base: string, result: string) {
     const httpAgent = new HttpAgent({ host: base });

packages/agent/src/agent/http/index.ts

@@ -352,12 +352,11 @@ export class HttpAgent implements Agent {
     return cbor.decode(await response.arrayBuffer());
   }
 
-  public async readState(
-    canisterId: Principal | string,
+  public async createReadStateRequest(
     fields: ReadStateOptions,
     identity?: Identity | Promise<Identity>,
-  ): Promise<ReadStateResponse> {
-    const canister = typeof canisterId === 'string' ? Principal.fromText(canisterId) : canisterId;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  ): Promise<any> {
     const id = await (identity !== undefined ? await identity : await this._identity);
     if (!id) {
       throw new IdentityInvalidError(
@@ -368,7 +367,7 @@ export class HttpAgent implements Agent {
 
     // TODO: remove this any. This can be a Signed or UnSigned request.
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    let transformedRequest: any = await this._transform({
+    const transformedRequest: any = await this._transform({
       request: {
         method: 'POST',
         headers: {
@@ -386,8 +385,19 @@ export class HttpAgent implements Agent {
     });
 
     // Apply transform for identity.
-    transformedRequest = await id?.transformRequest(transformedRequest);
+    return id?.transformRequest(transformedRequest);
+  }
+
+  public async readState(
+    canisterId: Principal | string,
+    fields: ReadStateOptions,
+    identity?: Identity | Promise<Identity>,
+    // eslint-disable-next-line
+    request?: any,
+  ): Promise<ReadStateResponse> {
+    const canister = typeof canisterId === 'string' ? Principal.fromText(canisterId) : canisterId;
 
+    const transformedRequest = request ?? (await this.createReadStateRequest(fields, identity));
     const body = cbor.encode(transformedRequest.body);
 
     const response = await this._fetch(

packages/agent/src/polling/index.ts

@@ -20,15 +20,19 @@ export type PollStrategyFactory = () => PollStrategy;
  * @param canisterId The effective canister ID.
  * @param requestId The Request ID to poll status for.
  * @param strategy A polling strategy.
+ * @param request Request for the readState call.
  */
 export async function pollForResponse(
   agent: Agent,
   canisterId: Principal,
   requestId: RequestId,
   strategy: PollStrategy,
+  // eslint-disable-next-line
+  request?: any,
 ): Promise<ArrayBuffer> {
   const path = [new TextEncoder().encode('request_status'), requestId];
-  const state = await agent.readState(canisterId, { paths: [path] });
+  const currentRequest = request ?? (await agent.createReadStateRequest?.({ paths: [path] }));
+  const state = await agent.readState(canisterId, { paths: [path] }, undefined, currentRequest);
   const cert = new Certificate(state, agent);
   const verified = await cert.verify();
   if (!verified) {
@@ -53,7 +57,7 @@ export async function pollForResponse(
     case RequestStatusResponseStatus.Processing:
       // Execute the polling strategy, then retry.
       await strategy(canisterId, requestId, status);
-      return pollForResponse(agent, canisterId, requestId, strategy);
+      return pollForResponse(agent, canisterId, requestId, strategy, currentRequest);
 
     case RequestStatusResponseStatus.Rejected: {
       const rejectCode = new Uint8Array(cert.lookup([...path, 'reject_code'])!)[0];