[Feature]: Types registry support (#117)

* Refactor entities module to support generic entity types and improve type safety

* Extend EntitiesModule with typed Todo handler for improved type safety in entities tests

* Enhance entity handling with typed sorting and import results

- Introduced `SortField` type for improved sorting parameter handling in list and filter methods.
- Updated `importEntities` method to return a structured `ImportResult` type, providing detailed import status and output.
- Modified list and filter methods to support generic field selection, enhancing type safety and flexibility.

* add support to augment the types of the sdk

* cleanups

* align comments

* update entities and function TypeRegistry in tests

* added EntityRecord

* added type registry for connectors

---------

Co-authored-by: Oz Sayag <ozs@wix.com>

Author: kfirstri

src/index.ts

@@ -36,6 +36,8 @@ export * from "./types.js";
 export type {
   EntitiesModule,
   EntityHandler,
+  EntityRecord,
+  EntityTypeRegistry,
   RealtimeEventType,
   RealtimeEvent,
   RealtimeCallback,
@@ -71,10 +73,14 @@ export type {
   CreateFileSignedUrlResult,
 } from "./modules/integrations.types.js";
 
-export type { FunctionsModule } from "./modules/functions.types.js";
+export type {
+  FunctionsModule,
+  FunctionNameRegistry,
+} from "./modules/functions.types.js";
 
 export type {
   AgentsModule,
+  AgentNameRegistry,
   AgentConversation,
   AgentMessage,
   AgentMessageReasoning,

src/modules/agents.types.ts

@@ -2,6 +2,19 @@ import { AxiosInstance } from "axios";
 import { RoomsSocket } from "../utils/socket-utils.js";
 import { ModelFilterParams } from "../types.js";
 
+/**
+ * Registry of agent names.
+ * Augment this interface to enable autocomplete for agent names.
+ */
+export interface AgentNameRegistry {}
+
+/**
+ * Agent name type - uses registry keys if augmented, otherwise string.
+ */
+export type AgentName = keyof AgentNameRegistry extends never
+  ? string
+  : keyof AgentNameRegistry;
+
 /**
  * Reasoning information for an agent message.
  *
@@ -135,7 +148,7 @@ export interface AgentMessage {
  */
 export interface CreateConversationParams {
   /** The name of the agent to create a conversation with. */
-  agent_name: string;
+  agent_name: AgentName;
   /** Optional metadata to attach to the conversation. */
   metadata?: Record<string, any>;
 }
@@ -376,5 +389,5 @@ When receiving messages through this function, tool call data is truncated for e
    * // User can open this URL to start a WhatsApp conversation
    * ```
    */
-  getWhatsAppConnectURL(agentName: string): string;
+  getWhatsAppConnectURL(agentName: AgentName): string;
 }

src/modules/connectors.types.ts

@@ -1,7 +1,16 @@
+/**
+ * Registry of connector integration types.
+ * Augment this interface to enable autocomplete for connector integration types.
+ */
+export interface ConnectorIntegrationTypeRegistry {}
+
 /**
  * The type of external integration/connector, such as `'googlecalendar'`, `'slack'`, or `'github'`.
+ * Uses registry keys if augmented, otherwise falls back to string.
  */
-export type ConnectorIntegrationType = string;
+export type ConnectorIntegrationType = keyof ConnectorIntegrationTypeRegistry extends never
+  ? string
+  : keyof ConnectorIntegrationTypeRegistry;
 
 /**
  * Response from the connectors access token endpoint.

src/modules/entities.types.ts

@@ -80,6 +80,37 @@ export type SortField<T> =
   | `+${keyof T & string}`
   | `-${keyof T & string}`;
 
+/**
+ * Fields added by the server to every entity record (id, dates, created_by, etc.).
+ */
+interface ServerEntityFields {
+  /** Unique identifier of the record */
+  id: string;
+  /** ISO 8601 timestamp when the record was created */
+  created_date: string;
+  /** ISO 8601 timestamp when the record was last updated */
+  updated_date: string;
+  /** Email of the user who created the record (may be hidden in some responses) */
+  created_by?: string | null;
+  /** ID of the user who created the record */
+  created_by_id?: string | null;
+  /** Whether the record is sample/seed data */
+  is_sample?: boolean;
+}
+
+/**
+ * Registry mapping entity names to their TypeScript types.
+ * Augment this interface with your entity schema (user-defined fields only).
+ */
+export interface EntityTypeRegistry {}
+
+/**
+ * Full record type for each entity: schema fields + server-injected fields (id, created_date, etc.).
+ */
+export type EntityRecord = {
+  [K in keyof EntityTypeRegistry]: EntityTypeRegistry[K] & ServerEntityFields;
+};
+
 /**
  * Entity handler providing CRUD operations for a specific entity type.
  *
@@ -303,7 +334,7 @@ export interface EntityHandler<T = any> {
    *   status: 'completed',
    *   priority: 'low'
    * });
-   * console.log('Deleted:', result);
+   * console.log('Deleted:', result.deleted);
    * ```
    */
   deleteMany(query: Partial<T>): Promise<DeleteManyResult>;
@@ -382,6 +413,20 @@ export interface EntityHandler<T = any> {
   subscribe(callback: RealtimeCallback<T>): () => void;
 }
 
+/**
+ * Typed entities module - maps registry keys to typed handlers (full record type).
+ */
+type TypedEntitiesModule = {
+  [K in keyof EntityTypeRegistry]: EntityHandler<EntityRecord[K]>;
+};
+
+/**
+ * Dynamic entities module - allows any entity name with untyped handler.
+ */
+type DynamicEntitiesModule = {
+  [entityName: string]: EntityHandler<any>;
+};
+
 /**
  * Entities module for managing app data.
  *
@@ -415,18 +460,4 @@ export interface EntityHandler<T = any> {
  * const allUsers = await base44.asServiceRole.entities.User.list();
  * ```
  */
-export interface EntitiesModule {
-  /**
-   * Access any entity by name.
-   *
-   * Use this to access entities defined in the app.
-   *
-   * @example
-   * ```typescript
-   * // Access entities dynamically
-   * base44.entities.MyEntity
-   * base44.entities.AnotherEntity
-   * ```
-   */
-  [entityName: string]: EntityHandler<any>;
-}
+export type EntitiesModule = TypedEntitiesModule & DynamicEntitiesModule;

src/modules/functions.types.ts

@@ -1,3 +1,16 @@
+/**
+ * Registry of function names.
+ * Augment this interface to enable autocomplete for function names.
+ */
+export interface FunctionNameRegistry {}
+
+/**
+ * Function name type - uses registry keys if augmented, otherwise string.
+ */
+export type FunctionName = keyof FunctionNameRegistry extends never
+  ? string
+  : keyof FunctionNameRegistry;
+
 /**
  * Functions module for invoking custom backend functions.
  *
@@ -46,5 +59,5 @@ export interface FunctionsModule {
    * };
    * ```
    */
-  invoke(functionName: string, data: Record<string, any>): Promise<any>;
+  invoke(functionName: FunctionName, data?: Record<string, any>): Promise<any>;
 }

tests/unit/entities.test.ts

@@ -12,10 +12,10 @@ interface Todo {
   completed: boolean;
 }
 
-// Declaration merging: extend EntitiesModule with typed Todo handler
+// Module augmentation: register Todo type in EntityTypeRegistry
 declare module "../../src/modules/entities.types.ts" {
-  interface EntitiesModule {
-    Todo: EntityHandler<Todo>;
+  interface EntityTypeRegistry {
+    Todo: Todo;
   }
 }
 

tests/unit/functions.test.ts

@@ -2,6 +2,15 @@ import { describe, test, expect, beforeEach, afterEach } from "vitest";
 import nock from "nock";
 import { createClient } from "../../src/index.ts";
 
+// Module augmentation: register function names in FunctionNameRegistry
+declare module "../../src/modules/functions.types.ts" {
+  interface FunctionNameRegistry {
+    sendNotification: true;
+    processOrder: true;
+    generateReport: true;
+  }
+}
+
 describe("Functions Module", () => {
   let base44: ReturnType<typeof createClient>;
   let scope;