From 6ef073dc6c7e0fcd2de8a55d1f244bc8bb2ad07d Mon Sep 17 00:00:00 2001
From: Andrew Bulat <andrii.bulat@gmail.com>
Date: Thu, 5 Dec 2024 06:20:38 +0000
Subject: [PATCH] Update `LiveMap.get` to always return `undefined` value as an
 option for keys that reference other LiveObjects

Referenced live objects can be not valid (haven't seen a create op) and
we should not surface such objects to the end user and return `undefined`
instead.
---
 ably.d.ts                                              | 6 +++---
 src/plugins/liveobjects/livemap.ts                     | 2 +-
 test/package/browser/template/src/ably.config.d.ts     | 4 ++--
 test/package/browser/template/src/index-liveobjects.ts | 4 +++-
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/ably.d.ts b/ably.d.ts
index 5d1aa1f82..938de0d7c 100644
--- a/ably.d.ts
+++ b/ably.d.ts
@@ -2103,12 +2103,12 @@ export type DefaultRoot =
  */
 export declare interface LiveMap<T extends LiveMapType> extends LiveObject<LiveMapUpdate> {
   /**
-   * Returns the value associated with a given key. Returns `undefined` if the key doesn't exist in a map.
+   * Returns the value associated with a given key. Returns `undefined` if the key doesn't exist in a map or if the associated {@link LiveObject} doesn't exist.
    *
    * @param key - The key to retrieve the value for.
-   * @returns A {@link LiveObject}, a primitive type (string, number, boolean, or binary data) or `undefined` if the key doesn't exist in a map.
+   * @returns A {@link LiveObject}, a primitive type (string, number, boolean, or binary data) or `undefined` if the key doesn't exist in a map or the associated {@link LiveObject} doesn't exist.
    */
-  get<TKey extends keyof T & string>(key: TKey): T[TKey];
+  get<TKey extends keyof T & string>(key: TKey): T[TKey] extends StateValue ? T[TKey] : T[TKey] | undefined;
 
   /**
    * Returns the number of key/value pairs in the map.
diff --git a/src/plugins/liveobjects/livemap.ts b/src/plugins/liveobjects/livemap.ts
index e4b1200ef..43eef94f6 100644
--- a/src/plugins/liveobjects/livemap.ts
+++ b/src/plugins/liveobjects/livemap.ts
@@ -85,7 +85,7 @@ export class LiveMap<T extends API.LiveMapType> extends LiveObject<LiveMapData,
    * - If the value is not an objectId, then that value is returned.
    */
   // force the key to be of type string as we only allow strings as key in a map
-  get<TKey extends keyof T & string>(key: TKey): T[TKey] {
+  get<TKey extends keyof T & string>(key: TKey): T[TKey] extends StateValue ? T[TKey] : T[TKey] | undefined {
     const element = this._dataRef.data.get(key);
 
     if (element === undefined) {
diff --git a/test/package/browser/template/src/ably.config.d.ts b/test/package/browser/template/src/ably.config.d.ts
index e5bca7718..3b3c69ddb 100644
--- a/test/package/browser/template/src/ably.config.d.ts
+++ b/test/package/browser/template/src/ably.config.d.ts
@@ -5,13 +5,13 @@ type CustomRoot = {
   stringKey: string;
   booleanKey: boolean;
   couldBeUndefined?: string;
-  mapKey?: LiveMap<{
+  mapKey: LiveMap<{
     foo: 'bar';
     nestedMap?: LiveMap<{
       baz: 'qux';
     }>;
   }>;
-  counterKey?: LiveCounter;
+  counterKey: LiveCounter;
 };
 
 declare global {
diff --git a/test/package/browser/template/src/index-liveobjects.ts b/test/package/browser/template/src/index-liveobjects.ts
index 1cd27b021..6bd7b3547 100644
--- a/test/package/browser/template/src/index-liveobjects.ts
+++ b/test/package/browser/template/src/index-liveobjects.ts
@@ -30,8 +30,10 @@ globalThis.testAblyPackage = async function () {
   const aBoolean: boolean = root.get('booleanKey');
   const couldBeUndefined: string | undefined = root.get('couldBeUndefined');
   // live objects on a root:
+  // LiveMap.get can still return undefined for LiveObject typed properties even if custom typings have them as non-optional.
+  // objects can be non-valid and result in the undefined value
   const counter: Ably.LiveCounter | undefined = root.get('counterKey');
-  const map: LiveObjectsTypes['root']['mapKey'] = root.get('mapKey');
+  const map: LiveObjectsTypes['root']['mapKey'] | undefined = root.get('mapKey');
   // check string literal types works
   // need to use nullish coalescing as we didn't actually create any data on the root,
   // so the next calls would fail. we only need to check that TypeScript types work