diff --git a/test/types.test.ts b/test/types.test.ts
index 6a1a38d..8a6b9a1 100644
--- a/test/types.test.ts
+++ b/test/types.test.ts
@@ -96,4 +96,12 @@ describe('Flagsmith Types', () => {
         // eslint-disable-next-line @typescript-eslint/no-unused-vars
         const incorrectNumberFlag: string = typedFlagsmith.getValue("numberFlag")
     });
+    test('should accept the same trait types in identify as in setTraits', async () => {
+        const { flagsmith } = getFlagsmith({});
+        const typedFlagsmith = flagsmith as IFlagsmith<string, "plan" | "session_id">
+        await typedFlagsmith.setTraits({ plan: "pro", session_id: { value: "abc", transient: true } })
+        await typedFlagsmith.identify("user", { plan: "pro", session_id: { value: "abc", transient: true } })
+        //@ts-expect-error - trait not defined
+        await typedFlagsmith.identify("user", { fail: "x" })
+    });
 });
diff --git a/types.d.ts b/types.d.ts
index 89c7816..63c6111 100644
--- a/types.d.ts
+++ b/types.d.ts
@@ -223,7 +223,7 @@ T extends string = string
     /**
      * Identify user, triggers a call to get flags if `flagsmith.init` has been called
      * */
-    identify: (userId: string, traits?: Record<T, IFlagsmithValue>) => Promise<void>;
+    identify: (userId: string, traits?: Record<T, IFlagsmithTrait>) => Promise<void>;
     /**
      * Retrieves the current state of flagsmith
      */