feat: add Agent Manager to utils package (#731)

# Motivation It is a simple but useful class to manage multiple HTTP agents associated with identities. # Credits The original scripts were created by @peterpeterparker in the [OISY repo](https://github.com/dfinity/oisy-wallet). This is just a transposition to a re-usable utility. # Changes New `AgentManager` class with: - `create` method: static constructor. - `createAgent` method: as the name suggests, it creates a new HTTP but do not cache it. - `getAgent` method: it fetches the HTTP agent among the cached, and if does note exists, it creates it and cache it. - `clearAgents` method: clean slate. # Tests I created a few tests for the few usecases. # Todos - [x] Maybe add more mocked identity, to check for concurrent multiple agents. --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
dfinity · Oct 24, 2024 · 6c79ecc · 6c79ecc
1 parent 7ed25cc
commit 6c79ecc
Show file tree

Hide file tree

Showing 6 changed files with 289 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@
 ## Features
 
 - Add support for `get_default_subnets` to `@dfinity/cmc`.
+- Add class `AgentManager` in `@dfinity/utils` which caches `HttpAgent` instances for different identities.
 
 # 2024.10.09-1140Z
 

diff --git a/packages/utils/README.md b/packages/utils/README.md
@@ -565,6 +565,70 @@ Parameters:
 
 [:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/services/canister.ts#L4)
 
+### :factory: AgentManager
+
+AgentManager class manages HttpAgent instances for different identities.
+
+It caches agents by identity to optimise resource usage and avoid unnecessary agent creation.
+Provides functionality to create new agents, retrieve cached agents, and clear the cache when needed.
+
+[:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/utils/agent.utils.ts#L53)
+
+#### Methods
+
+- [create](#gear-create)
+- [getAgent](#gear-getagent)
+- [clearAgents](#gear-clearagents)
+
+##### :gear: create
+
+Static factory method to create a new AgentManager instance.
+
+This method serves as an alternative to directly using the private constructor,
+making it more convenient to create instances of `AgentManager` using a simple and clear method.
+
+| Method   | Type                                           |
+| -------- | ---------------------------------------------- |
+| `create` | `(config: AgentManagerConfig) => AgentManager` |
+
+Parameters:
+
+- `config`: - Configuration options for the AgentManager instance.
+- `config.fetchRootKey`: - Whether to fetch the root key for certificate validation.
+- `config.host`: - The host to connect to.
+
+[:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/utils/agent.utils.ts#L69)
+
+##### :gear: getAgent
+
+Get or create an HTTP agent for a given identity.
+
+If the agent for the specified identity has been created and cached, it is retrieved from the cache.
+If no agent exists for the identity, a new one is created, cached, and then returned.
+
+| Method     | Type                                                             |
+| ---------- | ---------------------------------------------------------------- |
+| `getAgent` | `({ identity, }: { identity: Identity; }) => Promise<HttpAgent>` |
+
+Parameters:
+
+- `identity`: - The identity to be used to create the agent.
+
+[:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/utils/agent.utils.ts#L82)
+
+##### :gear: clearAgents
+
+Clear the cache of HTTP agents.
+
+This method removes all cached agents, forcing new agent creation on the next request for any identity.
+Useful when identities have changed or if you want to reset all active connections.
+
+| Method        | Type         |
+| ------------- | ------------ |
+| `clearAgents` | `() => void` |
+
+[:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/utils/agent.utils.ts#L114)
+
 ### :factory: InvalidPercentageError
 
 [:link: Source](https://github.com/dfinity/ic-js/tree/main/packages/utils/src/utils/asserts.utils.ts#L1)

diff --git a/packages/utils/src/mocks/agent.mock.ts b/packages/utils/src/mocks/agent.mock.ts
@@ -0,0 +1,39 @@
+import type { HttpAgent } from "@dfinity/agent";
+import { AgentManagerConfig } from "../utils/agent.utils";
+
+export const mockHttpAgent = {
+  call: jest.fn().mockResolvedValue({
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    json: async () => ({ data: "mocked call result" }),
+  }),
+  query: jest.fn().mockResolvedValue({
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    json: async () => ({ data: "mocked query result" }),
+  }),
+  fetchRootKey: jest.fn().mockResolvedValue(undefined),
+} as unknown as HttpAgent;
+
+export const mockHttpAgent2 = {
+  call: jest.fn().mockResolvedValue({
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    json: async () => ({ data: "mocked call result" }),
+  }),
+  query: jest.fn().mockResolvedValue({
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    json: async () => ({ data: "mocked query result" }),
+  }),
+  fetchRootKey: jest.fn().mockResolvedValue(undefined),
+} as unknown as HttpAgent;
+
+export const mockAgentManagerConfig: AgentManagerConfig = {
+  fetchRootKey: false,
+  host: "https://icp-api.io",
+};
diff --git a/packages/utils/src/mocks/identity.mock.ts b/packages/utils/src/mocks/identity.mock.ts
@@ -0,0 +1,20 @@
+import type { Identity } from "@dfinity/agent";
+import { Principal } from "@dfinity/principal";
+
+export const mockPrincipalText =
+  "xlmdg-vkosz-ceopx-7wtgu-g3xmd-koiyc-awqaq-7modz-zf6r6-364rh-oqe";
+
+export const mockPrincipal = Principal.fromText(mockPrincipalText);
+
+export const mockIdentity = {
+  getPrincipal: () => mockPrincipal,
+} as unknown as Identity;
+
+export const mockPrincipalText2 =
+  "5uuwe-ggtgm-fonrs-rblmx-cfc23-pb3dg-iyfk2-dle5w-j5uev-ggmep-6ae";
+
+export const mockPrincipal2 = Principal.fromText(mockPrincipalText2);
+
+export const mockIdentity2 = {
+  getPrincipal: () => mockPrincipal2,
+} as unknown as Identity;
diff --git a/packages/utils/src/utils/agent.utils.spec.ts b/packages/utils/src/utils/agent.utils.spec.ts
@@ -0,0 +1,86 @@
+import { HttpAgent } from "@dfinity/agent";
+import { describe, expect } from "@jest/globals";
+import {
+  mockAgentManagerConfig,
+  mockHttpAgent,
+  mockHttpAgent2,
+} from "../mocks/agent.mock";
+import { mockIdentity, mockIdentity2 } from "../mocks/identity.mock";
+import { AgentManager } from "./agent.utils";
+
+jest.mock("@dfinity/agent", () => ({
+  HttpAgent: {
+    create: jest.fn(),
+  },
+}));
+
+describe("AgentManager", () => {
+  let agentManager: AgentManager;
+  const mockHttpAgentCreate = HttpAgent.create as jest.Mock;
+
+  beforeEach(() => {
+    agentManager = AgentManager.create(mockAgentManagerConfig);
+    jest.clearAllMocks();
+
+    mockHttpAgentCreate.mockResolvedValueOnce(mockHttpAgent);
+  });
+
+  describe("getAgent", () => {
+    it("should create a new agent when there is none", async () => {
+      const agent = await agentManager.getAgent({ identity: mockIdentity });
+
+      expect(mockHttpAgentCreate).toHaveBeenCalledWith(
+        expect.objectContaining({ identity: mockIdentity }),
+      );
+      expect(agent).toBe(mockHttpAgent);
+    });
+
+    it("should return cached agent if already created", async () => {
+      await agentManager.getAgent({ identity: mockIdentity });
+
+      const agent = await agentManager.getAgent({ identity: mockIdentity });
+
+      expect(mockHttpAgentCreate).toHaveBeenCalledTimes(1);
+      expect(agent).toBe(mockHttpAgent);
+    });
+
+    it("should handle multiple agents for multiple identities", async () => {
+      await agentManager.getAgent({ identity: mockIdentity });
+
+      mockHttpAgentCreate.mockResolvedValueOnce(mockHttpAgent2);
+
+      const agent2 = await agentManager.getAgent({ identity: mockIdentity2 });
+      const agent1 = await agentManager.getAgent({ identity: mockIdentity });
+
+      expect(mockHttpAgentCreate).toHaveBeenCalledTimes(2);
+
+      expect(agent1).toBe(mockHttpAgent);
+      expect(agent1).not.toBe(mockHttpAgent2);
+
+      expect(agent2).toBe(mockHttpAgent2);
+      expect(agent2).not.toBe(mockHttpAgent);
+    });
+  });
+
+  describe("clearAgents", () => {
+    it("should clear cached agents", async () => {
+      await agentManager.getAgent({ identity: mockIdentity });
+
+      const agentBefore = await agentManager.getAgent({
+        identity: mockIdentity,
+      });
+
+      expect(agentBefore).toBe(mockHttpAgent);
+
+      agentManager.clearAgents();
+
+      const agentAfter = await agentManager.getAgent({
+        identity: mockIdentity,
+      });
+
+      expect(mockHttpAgentCreate).toHaveBeenCalledTimes(2);
+      expect(agentAfter).not.toBe(mockHttpAgent);
+      expect(agentAfter).toBeUndefined();
+    });
+  });
+});
diff --git a/packages/utils/src/utils/agent.utils.ts b/packages/utils/src/utils/agent.utils.ts
@@ -1,7 +1,7 @@
-import type { Agent } from "@dfinity/agent";
+import type { Agent, Identity } from "@dfinity/agent";
 import { AnonymousIdentity, HttpAgent } from "@dfinity/agent";
 import type { CreateAgentParams } from "../types/agent.utils";
-import { nonNullish } from "./nullish.utils";
+import { isNullish, nonNullish } from "./nullish.utils";
 
 /**
  * Get a default agent that connects to mainnet with the anonymous identity.
@@ -38,3 +38,80 @@ export const createAgent = async ({
     shouldFetchRootKey: fetchRootKey,
   });
 };
+
+export type AgentManagerConfig = Pick<
+  CreateAgentParams,
+  "fetchRootKey" | "host"
+>;
+
+/**
+ * AgentManager class manages HttpAgent instances for different identities.
+ *
+ * It caches agents by identity to optimise resource usage and avoid unnecessary agent creation.
+ * Provides functionality to create new agents, retrieve cached agents, and clear the cache when needed.
+ */
+export class AgentManager {
+  private agents: Record<string, HttpAgent> | undefined | null = undefined;
+
+  private constructor(private readonly config: AgentManagerConfig) {}
+
+  /**
+   * Static factory method to create a new AgentManager instance.
+   *
+   * This method serves as an alternative to directly using the private constructor,
+   * making it more convenient to create instances of `AgentManager` using a simple and clear method.
+   *
+   * @param {AgentManagerConfig} config - Configuration options for the AgentManager instance.
+   * @param {boolean} config.fetchRootKey - Whether to fetch the root key for certificate validation.
+   * @param {string} config.host - The host to connect to.
+   * @returns {AgentManager} A new instance of `AgentManager`.
+   */
+  public static create(config: AgentManagerConfig): AgentManager {
+    return new AgentManager(config);
+  }
+
+  /**
+   * Get or create an HTTP agent for a given identity.
+   *
+   * If the agent for the specified identity has been created and cached, it is retrieved from the cache.
+   * If no agent exists for the identity, a new one is created, cached, and then returned.
+   *
+   * @param {Identity} identity - The identity to be used to create the agent.
+   * @returns {Promise<HttpAgent>} The HttpAgent associated with the given identity.
+   */
+  public async getAgent({
+    identity,
+  }: {
+    identity: Identity;
+  }): Promise<HttpAgent> {
+    const key = identity.getPrincipal().toText();
+
+    if (isNullish(this.agents) || isNullish(this.agents[key])) {
+      const agent = await createAgent({
+        identity,
+        fetchRootKey: this.config.fetchRootKey,
+        host: this.config.host,
+        verifyQuerySignatures: true,
+      });
+
+      this.agents = {
+        ...(this.agents ?? {}),
+        [key]: agent,
+      };
+
+      return agent;
+    }
+
+    return this.agents[key];
+  }
+
+  /**
+   * Clear the cache of HTTP agents.
+   *
+   * This method removes all cached agents, forcing new agent creation on the next request for any identity.
+   * Useful when identities have changed or if you want to reset all active connections.
+   */
+  public clearAgents(): void {
+    this.agents = null;
+  }
+}