Bound what one guest can make the chat hold or compute (review, area 2)

bin/chat-chainlit.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 
 import chainlit as cl
+from chainlit.chat_context import chat_contexts
 from chainlit.context import context
 from chainlit.data.base import BaseDataLayer
 from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
@@ -168,6 +169,16 @@ async def resume(thread: ThreadDict) -> None:
 @cl.on_chat_end
 async def end() -> None:
     await static_messages(config, TriggerEvent.on_chat_end)
+    # Nothing of a finished session should outlive it in memory (review,
+    # area 2): Chainlit's own per-session message list, our offers, and -- for
+    # a guest, who can never resume -- the conversation's checkpoints.
+    sid = session_id()
+    chat_contexts.pop(sid, None)
+    proposals.drop_session(sid)
+    _seeding.pop(sid, None)
+    if cl.user_session.get("user") is None:
+        with contextlib.suppress(Exception):
+            await get_graph().forget_thread(current_thread_id())
 
 
 async def run_gsa_analysis(attachment: Attachment) -> None:
@@ -319,13 +330,17 @@ async def propose_gene_list(
         Proposal(
             # Declined, this goes to the model: name the genes when they came
             # from an earlier message, or "the list I gave you" means nothing.
+            # Bounded: each offer kept the whole message (up to 60K) and every
+            # identifier, ~0.85 MB, though only 3,000 can be submitted and the
+            # text only matters if the reader declines (review, area 2).
             text=(
-                f"{message.content}\n\n(The genes: {', '.join(identifiers)})"
+                f"{message.content[:MAX_MODEL_CHARS]}\n\n"
+                f"(The genes: {', '.join(identifiers[:50])})"
                 if earlier
-                else message.content
+                else message.content[:MAX_MODEL_CHARS]
             ),
             message_id=message.id,
-            identifiers=tuple(identifiers),
+            identifiers=tuple(identifiers[:MAX_SUBMITTED_IDENTIFIERS]),
             actions=((run.name, run.id), (decline.name, decline.id)),
         ),
     )
@@ -443,7 +458,10 @@ async def run_gene_list_analysis(text: str, identifiers: list[str]) -> None:
         seeded = await get_graph().seed_history(
             profile,
             thread_id=current_thread_id(),
-            messages=[HumanMessage(content=text), AIMessage(content=reply.for_model)],
+            messages=[
+                HumanMessage(content=text[:MAX_MODEL_CHARS]),
+                AIMessage(content=reply.for_model),
+            ],
         )
     except Exception:
         # The reader has their result; only follow-ups lose it.
@@ -550,6 +568,14 @@ async def answer_with_model(content: str, message_id: str) -> None:
     save_openai_metrics(message_id, openai_cb)
 
 
+#: What a question to the model may be. A gene list may be longer (it is
+#: read, not sent): `gene_list.MAX_MESSAGE_CHARS`.
+MAX_MODEL_CHARS = 8_000
+TOO_LONG = (
+    "That message is too long for me to answer. Please shorten it -- or, for "
+    "an expression matrix, attach it as a file with 📎."
+)
+
 #: Handoffs being seeded, per session. A message waits for its session's seed
 #: before touching the thread, or the seed lands beside a running turn and is
 #: lost while the chat says it is continuing from the summary.
@@ -592,6 +618,16 @@ def solve_key() -> str:
 
 @cl.on_message
 async def main(message: cl.Message) -> None:
+    try:
+        await handle_message(message)
+    finally:
+        # Chainlit appends every message to a per-session list before this
+        # runs, and nothing removes them: rejected ones too, kept until
+        # restart. Nothing here reads that list (review, area 2).
+        chat_contexts.pop(session_id(), None)
+
+
+async def handle_message(message: cl.Message) -> None:
     await wait_for_seed()
 
     # First, before any early return: a "yes" means the offer just made, so
@@ -604,6 +640,12 @@ async def main(message: cl.Message) -> None:
     if cl.user_session.get("user") is None and not _per_solve.allow(solve_key()):
         await cl.Message(content=PER_SOLVE_LIMITED).send()
         return
+    if len(message.content or "") > gene_list.MAX_MESSAGE_CHARS:
+        # Engine.io delivers up to a million characters. Each was kept in the
+        # thread's checkpoints and sent to the model on every later turn
+        # (review, area 2).
+        await cl.Message(content=TOO_LONG).send()
+        return
 
     await static_messages(config, TriggerEvent.on_message)
 
@@ -666,4 +708,9 @@ async def main(message: cl.Message) -> None:
         proposals.invite(session_id())
         return
 
+    if len(message.content or "") > MAX_MODEL_CHARS:
+        # Long enough to be a gene list, which has been handled above; too
+        # long for a question. The search page caps questions at 2,000.
+        await cl.Message(content=TOO_LONG).send()
+        return
     await answer_with_model(message.content, message.id)

src/analysis/gene_list.py

@@ -459,14 +459,20 @@ def describe_proposal(identifiers: list[str], *, earlier: bool = False) -> str:
 
 _CONFIRMS = re.compile(
     r"\A\s*(yes|y|yep|yeah|ok|okay|sure|go|go\s+ahead|run|run\s+it|do\s+it"
-    r"|please|please\s+do|yes,?\s+please|please\s+run\s+it)\s*[.!]*\s*\Z",
+    r"|please|please\s+do|yes,?\s+please|please\s+run\s+it)\s*+[.!]*+\s*+\Z",
     re.IGNORECASE,
 )
 
 
+#: No typed "yes" is longer than this. Checked first: the pattern ran on the
+#: raw message, and "y" plus a million spaces backtracked for an hour and a
+#: half on the event loop every session shares (review, area 2).
+MAX_CONFIRM_CHARS = 40
+
+
 def confirms(text: str) -> bool:
     """A typed yes to the proposal just made, instead of clicking Run."""
-    return bool(_CONFIRMS.match(text))
+    return len(text) <= MAX_CONFIRM_CHARS and bool(_CONFIRMS.match(text))
 
 
 @dataclass(frozen=True)

src/analysis/proposals.py

@@ -116,5 +116,9 @@ def take_invited(self, session_id: str) -> bool:
         invited, session.invited = session.invited, False
         return invited
 
+    def drop_session(self, session_id: str) -> None:
+        """Forget everything about a session that has ended."""
+        self._sessions.pop(session_id, None)
+
 
 proposals = ProposalStore()

tests/analysis/test_gene_list.py

@@ -615,3 +615,13 @@ def test_the_model_never_gets_the_link_that_carries_the_token() -> None:
     assert url in reply.text
     assert token not in reply.for_model
     assert "Regulation of TP53 Expression" in reply.for_model
+
+
+def test_a_padded_message_cannot_stall_the_yes_check() -> None:
+    # "y" plus spaces backtracked quadratically: ~90 minutes at 990K on the
+    # shared event loop (review, area 2). Bounded and possessive now.
+    started = time.perf_counter()
+    assert not confirms("y" + " " * 990_000 + "x")
+    assert not confirms("yes" + " " * 30 + "x")
+    assert time.perf_counter() - started < 0.1
+    assert confirms("yes  !")

tests/analysis/test_proposals.py

@@ -108,3 +108,14 @@ def test_the_last_list_is_kept_per_session() -> None:
     store.remember_list("s", ["EGFR", "KRAS"])
     assert store.last_list("s") == ["EGFR", "KRAS"]
     assert store.last_list("other") is None
+
+
+def test_a_finished_session_is_forgotten() -> None:
+    store = ProposalStore()
+    store.put("s", "p", offer())
+    store.invite("s")
+    store.remember_list("s", ["TP53", "MDM2"])
+    store.drop_session("s")
+    assert store.take("s", "p") is None
+    assert store.take_invited("s") is False
+    assert store.last_list("s") is None