From f6117a5fe3e065ce7b8ecfe9a91860e2637a64a2 Mon Sep 17 00:00:00 2001
From: Leon Zhao <leeeon233@gmail.com>
Date: Thu, 24 Apr 2025 16:21:48 +0800
Subject: [PATCH 1/3] fix: update EphemeralStore to support generic types

---
 crates/loro-wasm/index.ts                | 18 ++++++++---------
 crates/loro-wasm/tests/ephemeral.test.ts | 25 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/crates/loro-wasm/index.ts b/crates/loro-wasm/index.ts
index b3858acda..91a97c0b9 100644
--- a/crates/loro-wasm/index.ts
+++ b/crates/loro-wasm/index.ts
@@ -249,8 +249,8 @@ export class Awareness<T extends Value = Value> {
  * store2.apply(encoded);
  * ```
  */
-export class EphemeralStore<T extends Value = Value> {
-    inner: EphemeralStoreWasm<T>;
+export class EphemeralStore<T extends Record<string, Value> = Record<string, Value>> {
+    inner: EphemeralStoreWasm;
     private timer: number | undefined;
     private timeout: number;
     constructor(timeout: number = 30000) {
@@ -263,21 +263,21 @@ export class EphemeralStore<T extends Value = Value> {
         this.startTimerIfNotEmpty();
     }
 
-    set(key: string, value: T) {
-        this.inner.set(key, value);
+    set<K extends keyof T>(key: K, value: T[K]) {
+        this.inner.set(key as string, value);
         this.startTimerIfNotEmpty();
     }
 
-    get(key: string): T | undefined {
-        return this.inner.get(key);
+    get<K extends keyof T>(key: K): T[K] | undefined {
+        return this.inner.get(key as string);
     }
 
-    getAllStates(): Record<string, T> {
+    getAllStates(): Partial<T> {
         return this.inner.getAllStates();
     }
 
-    encode(key: string): Uint8Array {
-        return this.inner.encode(key);
+    encode<K extends keyof T>(key: K): Uint8Array {
+        return this.inner.encode(key as string);
     }
 
     encodeAll(): Uint8Array {
diff --git a/crates/loro-wasm/tests/ephemeral.test.ts b/crates/loro-wasm/tests/ephemeral.test.ts
index c7835116a..f0ece9043 100644
--- a/crates/loro-wasm/tests/ephemeral.test.ts
+++ b/crates/loro-wasm/tests/ephemeral.test.ts
@@ -116,6 +116,31 @@ describe("EphemeralStore", () => {
         });
     });
 
+    it("generic type", () => {
+        // Define a type to test type inference
+        const store = new EphemeralStore<{ foo: string, bar: number }>(10);
+        // This should compile correctly
+        store.set("foo", "bar");
+        store.set("bar", 1);
+
+        // @ts-expect-error - This should fail type checking as "foo" expects string
+        store.set("foo", 123);
+        // @ts-expect-error - This should fail type checking as "bar" expects number
+        store.set("bar", "string");
+        // @ts-expect-error - This should fail type checking as "baz" is not in the type
+        store.set("baz", "value");
+
+        // Verify runtime values are correct
+        expect(store.get("foo")).toBe("bar");
+        expect(store.get("bar")).toBe(1);
+
+        // Type inference for get should work too
+        const foo: string | undefined = store.get("foo");
+        const bar: number | undefined = store.get("bar");
+        expect(foo).toBe("bar");
+        expect(bar).toBe(1);
+    });
+
     it("subscribe", () => {
         const store = new EphemeralStore(10);
         let callTimes = 0;

From e8828f7eb9226519cfc4c37c2d59b2ce326eb34f Mon Sep 17 00:00:00 2001
From: Leon Zhao <leeeon233@gmail.com>
Date: Thu, 24 Apr 2025 16:31:21 +0800
Subject: [PATCH 2/3] fix: test

---
 crates/loro-wasm/tests/ephemeral.test.ts | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/crates/loro-wasm/tests/ephemeral.test.ts b/crates/loro-wasm/tests/ephemeral.test.ts
index f0ece9043..dd39fa6fa 100644
--- a/crates/loro-wasm/tests/ephemeral.test.ts
+++ b/crates/loro-wasm/tests/ephemeral.test.ts
@@ -123,13 +123,6 @@ describe("EphemeralStore", () => {
         store.set("foo", "bar");
         store.set("bar", 1);
 
-        // @ts-expect-error - This should fail type checking as "foo" expects string
-        store.set("foo", 123);
-        // @ts-expect-error - This should fail type checking as "bar" expects number
-        store.set("bar", "string");
-        // @ts-expect-error - This should fail type checking as "baz" is not in the type
-        store.set("baz", "value");
-
         // Verify runtime values are correct
         expect(store.get("foo")).toBe("bar");
         expect(store.get("bar")).toBe(1);
@@ -139,6 +132,13 @@ describe("EphemeralStore", () => {
         const bar: number | undefined = store.get("bar");
         expect(foo).toBe("bar");
         expect(bar).toBe(1);
+
+        // @ts-expect-error - This should fail type checking as "foo" expects string
+        store.set("foo", 123);
+        // @ts-expect-error - This should fail type checking as "bar" expects number
+        store.set("bar", "string");
+        // @ts-expect-error - This should fail type checking as "baz" is not in the type
+        store.set("baz", "value");
     });
 
     it("subscribe", () => {

From eb3a88452f1d9ae918c841aa144a558964f5eb52 Mon Sep 17 00:00:00 2001
From: Leon Zhao <leeeon233@gmail.com>
Date: Thu, 24 Apr 2025 17:29:11 +0800
Subject: [PATCH 3/3] fix: add delete for ephemeral

---
 crates/loro-wasm/index.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/crates/loro-wasm/index.ts b/crates/loro-wasm/index.ts
index 91a97c0b9..0b54c2500 100644
--- a/crates/loro-wasm/index.ts
+++ b/crates/loro-wasm/index.ts
@@ -268,6 +268,10 @@ export class EphemeralStore<T extends Record<string, Value> = Record<string, Val
         this.startTimerIfNotEmpty();
     }
 
+    delete<K extends keyof T>(key: K) {
+        this.inner.delete(key as string);
+    }
+
     get<K extends keyof T>(key: K): T[K] | undefined {
         return this.inner.get(key as string);
     }