From 44a22931f5e7b28d20815243a011f4d8ccfa041b Mon Sep 17 00:00:00 2001
From: Quinn Slack <quinn@slack.org>
Date: Sat, 12 Oct 2024 13:17:25 -0700
Subject: [PATCH] support forceHydration on PromptString (#5884)

PromptString instances can't be cloned because their reference value is
used as the key in the WeakMap that implements immutability and
encapsulation for PromptString.

(This broke something on a separate branch, so I extracted the fix and
test case to merge separately here.)

## Test plan

Added test case
---
 lib/shared/src/editor/hydrateAfterPostMessage.test.ts | 11 ++++++++++-
 lib/shared/src/editor/hydrateAfterPostMessage.ts      |  6 ++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/shared/src/editor/hydrateAfterPostMessage.test.ts b/lib/shared/src/editor/hydrateAfterPostMessage.test.ts
index 470c01c896d5..a5fb1948a134 100644
--- a/lib/shared/src/editor/hydrateAfterPostMessage.test.ts
+++ b/lib/shared/src/editor/hydrateAfterPostMessage.test.ts
@@ -1,6 +1,7 @@
 import { describe, expect, test } from 'vitest'
 
-import { hydrateAfterPostMessage, isDehydratedUri } from './hydrateAfterPostMessage'
+import { ps } from '../prompt/prompt-string'
+import { forceHydration, hydrateAfterPostMessage, isDehydratedUri } from './hydrateAfterPostMessage'
 
 // Mock URI hydration function
 const mockHydrateUri = (value: unknown) => {
@@ -115,3 +116,11 @@ describe('hydrateAfterPostMessage', () => {
         expect(originalValue.foo.bar).toEqual('baz')
     })
 })
+
+describe('forceHydration', () => {
+    test('handles PromptString', async () => {
+        const ps1 = ps`foo`
+        expect(ps1.toJSON()).toBe('foo')
+        expect(forceHydration(ps1).toJSON()).toBe('foo')
+    })
+})
diff --git a/lib/shared/src/editor/hydrateAfterPostMessage.ts b/lib/shared/src/editor/hydrateAfterPostMessage.ts
index 0444aada619f..9008c20bd028 100644
--- a/lib/shared/src/editor/hydrateAfterPostMessage.ts
+++ b/lib/shared/src/editor/hydrateAfterPostMessage.ts
@@ -1,5 +1,6 @@
 import type * as vscode from 'vscode'
 import type { URI } from 'vscode-uri'
+import { PromptString } from '../prompt/prompt-string'
 
 /**
  * Forces hydration by cloning any part that may be lazily hydrated. This is necessary before using
@@ -10,6 +11,11 @@ export function forceHydration(object: any): any {
     if (typeof object !== 'object' || object === null) {
         return object
     }
+    if (object instanceof PromptString) {
+        // Return as-is, because PromptString object references are used as keys in a WeakMap that
+        // implements immutability and encapsulation for PromptString.
+        return object
+    }
     if (Array.isArray(object)) {
         return object.map(forceHydration)
     }