From 4860a7e8578bc10bcca4ccef70f51ae3ae06ac77 Mon Sep 17 00:00:00 2001
From: Tai Sakuma <tai.sakuma@gmail.com>
Date: Mon, 13 Nov 2023 13:49:29 -0500
Subject: [PATCH] Yield cached stdout before subscribing to new output

---
 nextlinegraphql/plugins/ctrl/__init__.py      | 30 ++++++++++++++++++-
 .../plugins/ctrl/schema/subscription.py       |  2 ++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/nextlinegraphql/plugins/ctrl/__init__.py b/nextlinegraphql/plugins/ctrl/__init__.py
index f2659b9..810aea6 100644
--- a/nextlinegraphql/plugins/ctrl/__init__.py
+++ b/nextlinegraphql/plugins/ctrl/__init__.py
@@ -1,10 +1,13 @@
 __all__ = ['Plugin']
 
+import asyncio
 from collections.abc import AsyncIterator, MutableMapping
 
 import apluggy as pluggy
 from apluggy import asynccontextmanager
 from nextline import Nextline
+from nextline.types import StdoutInfo
+from nextline.utils import merge_aiters
 
 from nextlinegraphql.hook import spec
 
@@ -24,15 +27,40 @@ async def update_lifespan_context(
         run_no: int = max(hook.hook.initial_run_no(), default=1)
         script: str = [*hook.hook.initial_script(), statement][0]
         self._nextline = Nextline(script, run_no)
+        self._stdout_cache = list[str]()
         context['nextline'] = self._nextline
 
     @spec.hookimpl(trylast=True)  # trylast so to be the innermost context
     @asynccontextmanager
     async def lifespan(self) -> AsyncIterator[None]:
         '''Yield within the nextline context.'''
-        async with self._nextline:
+        async with (self._nextline, self._cache_stdout()):
             yield
 
     @spec.hookimpl
     def update_strawberry_context(self, context: MutableMapping) -> None:
         context['nextline'] = self._nextline
+        context['stdout_cache'] = self._stdout_cache
+
+    @asynccontextmanager
+    async def _cache_stdout(self) -> AsyncIterator[None]:
+        task = asyncio.create_task(self.__cache_stdout())
+        try:
+            yield
+        finally:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+    async def __cache_stdout(self) -> None:
+        aiters = merge_aiters(
+            self._nextline.subscribe_state(), self._nextline.subscribe_stdout()
+        )
+        async for _, v in aiters:
+            match v:
+                case 'initialized':
+                    self._stdout_cache.clear()
+                case StdoutInfo(text=text) if text is not None:
+                    self._stdout_cache.append(text)
diff --git a/nextlinegraphql/plugins/ctrl/schema/subscription.py b/nextlinegraphql/plugins/ctrl/schema/subscription.py
index 07707f8..d43b9f1 100644
--- a/nextlinegraphql/plugins/ctrl/schema/subscription.py
+++ b/nextlinegraphql/plugins/ctrl/schema/subscription.py
@@ -59,6 +59,8 @@ async def subscribe_prompting(
 
 async def subscribe_stdout(info: Info) -> AsyncIterator[str]:
     nextline: Nextline = info.context["nextline"]
+    stdout_cache: list[str] = info.context["stdout_cache"]
+    yield ''.join(stdout_cache)
     async for i in nextline.subscribe_stdout():
         assert i.text is not None
         yield i.text