fix(jsx): fix context provider with async component (#2124)

* test(jsx): add test for Context with Suspense * fix(jsx): fix context provider with async component * chore: denoify
honojs · Jan 31, 2024 · 0ed91a7 · 0ed91a7
1 parent f970a64
commit 0ed91a7
Show file tree

Hide file tree

Showing 3 changed files with 241 additions and 34 deletions.
diff --git a/deno_dist/jsx/index.ts b/deno_dist/jsx/index.ts
@@ -88,12 +88,14 @@ const childrenToStringToBuffer = (children: Child[], buffer: StringBuffer): void
   }
 }
 
+type LocalContexts = [Context<unknown>, unknown][]
 export type Child = string | Promise<string> | number | JSXNode | Child[]
 export class JSXNode implements HtmlEscaped {
   tag: string | Function
   props: Props
   children: Child[]
   isEscaped: true = true as const
+  localContexts?: LocalContexts
   constructor(tag: string | Function, props: Props, children: Child[]) {
     this.tag = tag
     this.props = props
@@ -102,7 +104,16 @@ export class JSXNode implements HtmlEscaped {
 
   toString(): string | Promise<string> {
     const buffer: StringBuffer = ['']
-    this.toStringToBuffer(buffer)
+    this.localContexts?.forEach(([context, value]) => {
+      context.values.push(value)
+    })
+    try {
+      this.toStringToBuffer(buffer)
+    } finally {
+      this.localContexts?.forEach(([context]) => {
+        context.values.pop()
+      })
+    }
     return buffer.length === 1 ? buffer[0] : stringBufferToString(buffer)
   }
 
@@ -178,7 +189,24 @@ class JSXFunctionNode extends JSXNode {
     })
 
     if (res instanceof Promise) {
-      buffer.unshift('', res)
+      if (globalContexts.length === 0) {
+        buffer.unshift('', res)
+      } else {
+        // save current contexts for resuming
+        const currentContexts: LocalContexts = globalContexts.map((c) => [
+          c,
+          c.values[c.values.length - 1],
+        ])
+        buffer.unshift(
+          '',
+          res.then((childRes) => {
+            if (childRes instanceof JSXNode) {
+              childRes.localContexts = currentContexts
+            }
+            return childRes
+          })
+        )
+      }
     } else if (res instanceof JSXNode) {
       res.toStringToBuffer(buffer)
     } else if (typeof res === 'number' || (res as HtmlEscaped).isEscaped) {
@@ -259,33 +287,37 @@ export interface Context<T> {
   Provider: FC<{ value: T }>
 }
 
+const globalContexts: Context<unknown>[] = []
+
 export const createContext = <T>(defaultValue: T): Context<T> => {
   const values = [defaultValue]
-  return {
+  const context: Context<T> = {
     values,
     Provider(props): HtmlEscapedString | Promise<HtmlEscapedString> {
       values.push(props.value)
-      const string = props.children
-        ? (Array.isArray(props.children)
-            ? new JSXFragmentNode('', {}, props.children)
-            : props.children
-          ).toString()
-        : ''
-      values.pop()
+      let string
+      try {
+        string = props.children
+          ? (Array.isArray(props.children)
+              ? new JSXFragmentNode('', {}, props.children)
+              : props.children
+            ).toString()
+          : ''
+      } finally {
+        values.pop()
+      }
 
       if (string instanceof Promise) {
-        return Promise.resolve().then<HtmlEscapedString>(async () => {
-          values.push(props.value)
-          const awaited = await string
-          const promiseRes = raw(awaited, (awaited as HtmlEscapedString).callbacks)
-          values.pop()
-          return promiseRes
-        })
+        return string.then((resString) =>
+          raw(resString, (resString as HtmlEscapedString).callbacks)
+        )
       } else {
         return raw(string)
       }
     },
   }
+  globalContexts.push(context as Context<unknown>)
+  return context
 }
 
 export const useContext = <T>(context: Context<T>): T => {

diff --git a/src/jsx/index.test.tsx b/src/jsx/index.test.tsx
@@ -2,6 +2,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { html } from '../helper/html'
 import { Hono } from '../hono'
+import { Suspense, renderToReadableStream } from './streaming'
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 import { jsx, memo, Fragment, createContext, useContext } from './index'
 import type { Context, FC } from './index'
@@ -491,12 +492,25 @@ describe('Fragment', () => {
 describe('Context', () => {
   let ThemeContext: Context<string>
   let Consumer: FC
+  let ErrorConsumer: FC
+  let AsyncConsumer: FC
+  let AsyncErrorConsumer: FC
   beforeAll(() => {
     ThemeContext = createContext('light')
     Consumer = () => {
       const theme = useContext(ThemeContext)
       return <span>{theme}</span>
     }
+    ErrorConsumer = () => {
+      throw new Error('ErrorConsumer')
+    }
+    AsyncConsumer = async () => {
+      const theme = useContext(ThemeContext)
+      return <span>{theme}</span>
+    }
+    AsyncErrorConsumer = async () => {
+      throw new Error('AsyncErrorConsumer')
+    }
   })
 
   describe('with .Provider', () => {
@@ -535,10 +549,139 @@ describe('Context', () => {
       )
       expect(template.toString()).toBe('<span>dark</span><span>black</span><span>dark</span>')
     })
+
+    it('should reset context by error', () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <ErrorConsumer />
+        </ThemeContext.Provider>
+      )
+      expect(() => template.toString()).toThrow()
+
+      const nextRequest = <Consumer />
+      expect(nextRequest.toString()).toBe('<span>light</span>')
+    })
   })
 
   it('default value', () => {
     const template = <Consumer />
     expect(template.toString()).toBe('<span>light</span>')
   })
+
+  describe('with Suspence', () => {
+    const RedTheme = () => (
+      <ThemeContext.Provider value='red'>
+        <Consumer />
+      </ThemeContext.Provider>
+    )
+
+    it('Should preserve context in sync component', async () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <Suspense fallback={<RedTheme />}>
+            <Consumer />
+            <ThemeContext.Provider value='black'>
+              <Consumer />
+            </ThemeContext.Provider>
+          </Suspense>
+        </ThemeContext.Provider>
+      )
+      const stream = renderToReadableStream(template)
+
+      const chunks = []
+      const textDecoder = new TextDecoder()
+      for await (const chunk of stream as any) {
+        chunks.push(textDecoder.decode(chunk))
+      }
+
+      expect(chunks).toEqual(['<span>dark</span><span>black</span>'])
+    })
+
+    it('Should preserve context in async component', async () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <Suspense fallback={<RedTheme />}>
+            <Consumer />
+            <ThemeContext.Provider value='black'>
+              <AsyncConsumer />
+            </ThemeContext.Provider>
+          </Suspense>
+        </ThemeContext.Provider>
+      )
+      const stream = renderToReadableStream(template)
+
+      const chunks = []
+      const textDecoder = new TextDecoder()
+      for await (const chunk of stream as any) {
+        chunks.push(textDecoder.decode(chunk))
+      }
+
+      expect(chunks).toEqual([
+        '<template id="H:0"></template><span>red</span><!--/$-->',
+        `<template><span>dark</span><span>black</span></template><script>
+((d,c,n) => {
+c=d.currentScript.previousSibling
+d=d.getElementById('H:0')
+if(!d)return
+do{n=d.nextSibling;n.remove()}while(n.nodeType!=8||n.nodeValue!='/$')
+d.replaceWith(c.content)
+})(document)
+</script>`,
+      ])
+    })
+  })
+
+  describe('async component', () => {
+    const ParentAsyncConsumer = async () => {
+      const theme = useContext(ThemeContext)
+      return (
+        <div>
+          <span>{theme}</span>
+          <AsyncConsumer />
+        </div>
+      )
+    }
+
+    const ParentAsyncErrorConsumer = async () => {
+      const theme = useContext(ThemeContext)
+      return (
+        <div>
+          <span>{theme}</span>
+          <AsyncErrorConsumer />
+        </div>
+      )
+    }
+
+    it('simple', async () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <AsyncConsumer />
+        </ThemeContext.Provider>
+      )
+      expect((await template.toString()).toString()).toBe('<span>dark</span>')
+    })
+
+    it('nested', async () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <ParentAsyncConsumer />
+        </ThemeContext.Provider>
+      )
+      expect((await template.toString()).toString()).toBe(
+        '<div><span>dark</span><span>dark</span></div>'
+      )
+    })
+
+    it('should reset context by error', async () => {
+      const template = (
+        <ThemeContext.Provider value='dark'>
+          <ParentAsyncErrorConsumer />
+        </ThemeContext.Provider>
+      )
+      expect(async () => (await template.toString()).toString()).rejects.toThrow()
+
+      const nextRequest = <Consumer />
+      expect(nextRequest.toString()).toBe('<span>light</span>')
+    })
+  })
 })
diff --git a/src/jsx/index.ts b/src/jsx/index.ts
@@ -88,12 +88,14 @@ const childrenToStringToBuffer = (children: Child[], buffer: StringBuffer): void
   }
 }
 
+type LocalContexts = [Context<unknown>, unknown][]
 export type Child = string | Promise<string> | number | JSXNode | Child[]
 export class JSXNode implements HtmlEscaped {
   tag: string | Function
   props: Props
   children: Child[]
   isEscaped: true = true as const
+  localContexts?: LocalContexts
   constructor(tag: string | Function, props: Props, children: Child[]) {
     this.tag = tag
     this.props = props
@@ -102,7 +104,16 @@ export class JSXNode implements HtmlEscaped {
 
   toString(): string | Promise<string> {
     const buffer: StringBuffer = ['']
-    this.toStringToBuffer(buffer)
+    this.localContexts?.forEach(([context, value]) => {
+      context.values.push(value)
+    })
+    try {
+      this.toStringToBuffer(buffer)
+    } finally {
+      this.localContexts?.forEach(([context]) => {
+        context.values.pop()
+      })
+    }
     return buffer.length === 1 ? buffer[0] : stringBufferToString(buffer)
   }
 
@@ -178,7 +189,24 @@ class JSXFunctionNode extends JSXNode {
     })
 
     if (res instanceof Promise) {
-      buffer.unshift('', res)
+      if (globalContexts.length === 0) {
+        buffer.unshift('', res)
+      } else {
+        // save current contexts for resuming
+        const currentContexts: LocalContexts = globalContexts.map((c) => [
+          c,
+          c.values[c.values.length - 1],
+        ])
+        buffer.unshift(
+          '',
+          res.then((childRes) => {
+            if (childRes instanceof JSXNode) {
+              childRes.localContexts = currentContexts
+            }
+            return childRes
+          })
+        )
+      }
     } else if (res instanceof JSXNode) {
       res.toStringToBuffer(buffer)
     } else if (typeof res === 'number' || (res as HtmlEscaped).isEscaped) {
@@ -259,33 +287,37 @@ export interface Context<T> {
   Provider: FC<{ value: T }>
 }
 
+const globalContexts: Context<unknown>[] = []
+
 export const createContext = <T>(defaultValue: T): Context<T> => {
   const values = [defaultValue]
-  return {
+  const context: Context<T> = {
     values,
     Provider(props): HtmlEscapedString | Promise<HtmlEscapedString> {
       values.push(props.value)
-      const string = props.children
-        ? (Array.isArray(props.children)
-            ? new JSXFragmentNode('', {}, props.children)
-            : props.children
-          ).toString()
-        : ''
-      values.pop()
+      let string
+      try {
+        string = props.children
+          ? (Array.isArray(props.children)
+              ? new JSXFragmentNode('', {}, props.children)
+              : props.children
+            ).toString()
+          : ''
+      } finally {
+        values.pop()
+      }
 
       if (string instanceof Promise) {
-        return Promise.resolve().then<HtmlEscapedString>(async () => {
-          values.push(props.value)
-          const awaited = await string
-          const promiseRes = raw(awaited, (awaited as HtmlEscapedString).callbacks)
-          values.pop()
-          return promiseRes
-        })
+        return string.then((resString) =>
+          raw(resString, (resString as HtmlEscapedString).callbacks)
+        )
       } else {
         return raw(string)
       }
     },
   }
+  globalContexts.push(context as Context<unknown>)
+  return context
 }
 
 export const useContext = <T>(context: Context<T>): T => {