feat: improve typings with Constructor type

Author: priley86

src/index.ts

@@ -1,14 +1,20 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { ApiClient as Auth0APIClient } from "@auth0/auth0-api-js";
+import type { AIChatAgent } from "agents/ai-chat-agent";
 import { AsyncLocalStorage } from "node:async_hooks";
-import { Connection, ConnectionContext, WSMessage } from "partyserver";
+import { Connection, ConnectionContext, Server, WSMessage } from "partyserver";
 import {
   InsufficientScopeError,
   InvalidTokenError,
   UnauthorizedError,
 } from "./bearer/errors.js";
 import getToken from "./bearer/index.js";
-import { AuthenticatedServer, TokenSet, WithAuthParams } from "./types.js";
+import {
+  AuthenticatedServer,
+  Constructor,
+  TokenSet,
+  WithAuthParams,
+} from "./types.js";
 
 export interface Token {
   sub?: string;
@@ -56,11 +62,11 @@ function validateScopes(
  */
 export const WithAuth = <
   Env extends { AUTH0_DOMAIN: string; AUTH0_AUDIENCE: string },
-  TBase extends new (...args: any[]) => any,
+  TBase extends Constructor<Server<Env>> | Constructor<AIChatAgent<Env>>,
 >(
   Base: TBase,
   options: WithAuthParams = { authRequired: true },
-): TBase & (new (...args: any[]) => AuthenticatedServer) => {
+) => {
   const authRequired = options.authRequired ?? true;
   const debug = options.debug ?? (() => {});
   // I had to do this because:
@@ -326,7 +332,7 @@ export const WithAuth = <
       tokenSetPerConnection.delete(connection.id);
       super.onClose(connection, code, reason, wasClean);
     }
-  } as TBase & (new (...args: any[]) => AuthenticatedServer);
+  };
 };
 
 export { OwnedAgent, WithOwnership } from "./withOwnership.js";

src/withOwnership.ts

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import { Connection } from "partyserver";
-import { AuthorizedServer } from "./types.js";
+import type { AIChatAgent } from "agents/ai-chat-agent";
+import { Connection, Server } from "partyserver";
+import { AuthenticatedServer, AuthorizedServer, Constructor } from "./types.js";
 
 interface DurableObjectState {
   readonly storage: {
@@ -25,17 +26,21 @@ const isDurableObjectState = (state: any): state is DurableObjectState => {
  *
  * Every time a connection or request is made, the ownership is checked.
  *
- * @param Base - The base class to extend from.
+ * @param Base - The base class to extend from. This should be a class that extends `Server` or `AIChatAgent` and implements `AuthenticatedServer`.
  *
  * @returns - A new class that extends the base class with ownership functionality.
  */
-export const WithOwnership = <TBase extends new (...args: any[]) => any>(
+export const WithOwnership = <
+  TBase extends
+    | Constructor<Server<any> & AuthenticatedServer>
+    | Constructor<AIChatAgent<any> & AuthenticatedServer>,
+>(
   Base: TBase,
   options: { debug?: (message: string, ctx: any) => void } = {},
-): TBase & (new (...args: any[]) => AuthorizedServer) => {
+): TBase & Constructor<AuthorizedServer> => {
   const debug = options.debug ?? (() => {});
 
-  return class extends Base {
+  return class WithOwnership extends Base {
     async onAuthorizedConnect(connection: any, ctx: any): Promise<void> {
       debug("Authenticated connection", {
         connection,
@@ -58,7 +63,8 @@ export const WithOwnership = <TBase extends new (...args: any[]) => any>(
      * @returns - A boolean indicating if the current user is the owner.
      */
     async #isCurrentUserOwner(): Promise<boolean> {
-      const userInfo = this.getClaims();
+      // Base class implements AuthenticatedServer, so getClaims() is available
+      const userInfo = (this as any as AuthenticatedServer).getClaims();
       const objectOwner = await this.getOwner();
       if (objectOwner !== userInfo?.sub) {
         return false;
@@ -87,6 +93,8 @@ export const WithOwnership = <TBase extends new (...args: any[]) => any>(
     }
 
     #getDurableStorage() {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-expect-error
       const ctx = this.ctx;
       if (!isDurableObjectState(ctx)) {
         throw new Error(
@@ -110,7 +118,7 @@ export const WithOwnership = <TBase extends new (...args: any[]) => any>(
     async getOwner(): Promise<string | undefined> {
       return this.#getDurableStorage().get("owner");
     }
-  } as TBase & (new (...args: any[]) => AuthorizedServer);
+  } as unknown as TBase & Constructor<AuthorizedServer>;
 };
 
 /**

test/flumix-compatibility.test.ts

@@ -45,7 +45,7 @@ describe("Flumix Compatibility", () => {
     // Test composing multiple mixins
     type BaseClass = typeof AIChatAgent<TestEnv>;
     type WithAuthClass = ReturnType<typeof AuthAgent<TestEnv, BaseClass>>;
-    type ComposedClass = ReturnType<typeof OwnedAgent<WithAuthClass>>;
+    type ComposedClass = ReturnType<typeof OwnedAgent<TestEnv, WithAuthClass>>;
     type ComposedInstance = InstanceType<ComposedClass>;
 
     // Verify it has methods from both mixins
@@ -80,15 +80,39 @@ describe("Flumix Compatibility", () => {
     // import it, triggering the cloudflare: protocol error. Instead, we verify
     // that our mixins work with flumix using a mock base class.
 
+    // Mock the AIChatAgent class structure for type compatibility
     class MockServer<Env> {
       constructor(
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        public party: any,
+        public ctx: any,
         public env?: Env,
       ) {}
+      
+      // Add stub methods to satisfy AIChatAgent interface
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      async fetch(_request: Request): Promise<Response> {
+        return new Response();
+      }
+      
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars
+      async onConnect(_connection: any, _ctx: any): Promise<void> {}
+      
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars
+      async onMessage(_connection: any, _message: any): Promise<void> {}
+      
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      async onRequest(_request: Request): Promise<Response> {
+        return new Response();
+      }
+      
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars
+      async onClose(_connection: any, _code: number, _reason: string, _wasClean: boolean): Promise<void> {}
+      
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars
+      onError(_connection: any, _error: unknown): void {}
     }
 
-    const SuperAgent = extend(MockServer<TestEnv>)
+    const SuperAgent = extend(MockServer<TestEnv> as unknown as typeof AIChatAgent<TestEnv>)
       .with(AuthAgent)
       .with(OwnedAgent)
       .build();
@@ -123,8 +147,8 @@ describe("Flumix Compatibility", () => {
       AUTH0_DOMAIN: "test.auth0.com",
       AUTH0_AUDIENCE: "https://api.test.com",
     };
-    const mockParty = {};
-    const instance = new SuperAgent(mockParty, mockEnv);
+    const mockCtx = {};
+    const instance = new SuperAgent(mockCtx, mockEnv);
     expect(instance).toBeDefined();
     expect(typeof instance.getClaims).toBe("function");
     // eslint-disable-next-line @typescript-eslint/no-explicit-any