diff --git a/memory/message_memory.py b/memory/message_memory.py index 46ae60ff..426b020b 100644 --- a/memory/message_memory.py +++ b/memory/message_memory.py @@ -11,7 +11,6 @@ ResponseExtractor, ) - async def safe_call( call, *args, @@ -28,6 +27,39 @@ async def safe_call( ) +async def emit_runtime_memory_update( + context, +) -> None: + + emitter = getattr( + context, + "emitter", + None, + ) + emit = getattr( + emitter, + "emit", + None, + ) + + await safe_call( + emit, + { + "type": "runtime_memory_update", + "memory": getattr( + context, + "runtime_memory", + "", + ), + "updates": getattr( + context, + "runtime_memory_updates", + 0, + ), + }, + ) + + def build_runtime_memory_system_prompt() -> str: return ( @@ -36,19 +68,33 @@ def build_runtime_memory_system_prompt() -> str: "Do not output JSON.\n" "Do not use Markdown headings.\n" "Do not explain your reasoning or the summarization process.\n" + "Write memory as atomic bullet lines, one semantic entity per line.\n" + "Each line should start with a compact semantic label such as topic, " + "user intent, active topic, open reference, pending choice, offered options, " + "preference, decision, pattern, or interrupted response.\n" + "Avoid writing about JIN's role unless the role itself changed or matters. " + "Describe assistant actions neutrally instead.\n" + "Keep memory actionable: write what helps the next answer, not a recap of " + "what happened.\n" + "Do not merge unrelated facts into one sentence. Prefer separate lines " + "over broad phrasing like 'Topic established: X, specifically Y'.\n" + "Finish every bullet line completely. Never leave a line mid-phrase.\n" "Preserve still-relevant existing memory. Update it instead of replacing it blindly.\n" "Drop old details only when they are clearly obsolete, duplicated, or no longer useful.\n" "Decide the summary depth from the signal in the latest turn.\n" "Use shallow summarization for simple factual, isolated, or low-signal turns: " - "keep only the dry fact, topic, or unresolved reference that could help " - "the next answer.\n" + "keep one or two bullet lines with only the dry fact, topic, or unresolved " + "reference that could help the next answer.\n" "Use deep summarization for turns that reveal user intent, project direction, " "preferences, recurring patterns, decisions, emotional tone, or a meaningful " - "shift in the conversation trajectory.\n" + "shift in the conversation trajectory; use three to six bullet lines when " + "the turn carries that much signal.\n" "If the user asks a follow-up that depends on prior context, preserve the " "referent clearly enough for the next brain prompt to resolve it.\n" "If the user switches topic, keep the new topic without forcing it into the " "previous one; only note a pattern if the sequence itself is meaningful.\n" + "If the assistant response was aborted or incomplete, mark it as incomplete " + "and do not treat it as resolved.\n" "Do not infer stable user traits from a single turn.\n" "Do not over-interpret jokes, tests, or casual topic changes.\n" "Prefer compact continuity over transcript-like detail.\n" @@ -73,7 +119,30 @@ def build_runtime_memory_user_prompt( f"{user_message.strip()}\n\n" "Latest JIN answer:\n" f"{assistant_message.strip()}\n\n" - "Rewrite the runtime memory now." + "Rewrite the runtime memory now as atomic bullet lines." + ) + + +def build_interrupted_assistant_message( + *, + user_message: str, + assistant_message: str, +) -> str: + + partial_text = assistant_message.strip() + + if not partial_text: + partial_text = ( + "No complete assistant answer was delivered." + ) + + return ( + "Assistant response was interrupted by the user and is incomplete. " + "Do not treat this turn as resolved.\n\n" + "Interrupted user topic/request:\n" + f"{user_message.strip()}\n\n" + "Partial assistant text before interruption:\n" + f"{partial_text}" ) @@ -93,6 +162,97 @@ def extract_runtime_memory_text( return text.strip() +def is_runtime_memory_response_truncated( + response: dict, +) -> bool: + + finish_reason = ( + ResponseExtractor + .extract_finish_reason( + response + ) + .lower() + ) + + return finish_reason in ( + "length", + "max_tokens", + ) + + +def looks_like_incomplete_runtime_memory( + text: str, +) -> bool: + + stripped = ( + text + or "" + ).strip() + + if not stripped: + return True + + if stripped[-1] in ( + ",", + ":", + "(", + "[", + "{", + ): + return True + + pairs = ( + ( + "(", + ")", + ), + ( + "[", + "]", + ), + ( + "{", + "}", + ), + ) + + return any( + stripped.count(open_char) + > stripped.count(close_char) + for open_char, close_char + in pairs + ) + + +async def ask_runtime_memory_model( + *, + service_client, + current_memory: str, + user_message: str, + assistant_message: str, +) -> dict: + + return await ask_service_model( + client=service_client, + system_prompt=( + build_runtime_memory_system_prompt() + ), + user_prompt=( + build_runtime_memory_user_prompt( + current_memory=current_memory, + user_message=user_message, + assistant_message=assistant_message, + ) + ), + temperature=( + config.SERVICE_TEMPERATURE + ), + max_tokens=( + config.SERVICE_MAX_TOKENS + ), + ) + + async def summarize_runtime_memory( *, context, @@ -132,33 +292,43 @@ async def summarize_runtime_memory( ) try: - response = await ask_service_model( - client=service_client, - system_prompt=( - build_runtime_memory_system_prompt() - ), - user_prompt=( - build_runtime_memory_user_prompt( - current_memory=current_memory, - user_message=user_message, - assistant_message=assistant_message, - ) - ), - temperature=( - config.SERVICE_TEMPERATURE - ), - max_tokens=( - min( - config.SERVICE_MAX_TOKENS, - 512, - ) - ), + response = await ask_runtime_memory_model( + service_client=service_client, + current_memory=current_memory, + user_message=user_message, + assistant_message=assistant_message, ) updated_memory = extract_runtime_memory_text( response ) + if ( + is_runtime_memory_response_truncated( + response + ) + or looks_like_incomplete_runtime_memory( + updated_memory + ) + ): + await safe_call( + getattr( + getattr( + context, + "logger", + None, + ), + "log_error", + None, + ), + "[MEMORY] runtime memory update skipped", + details=( + "Summarizer returned an incomplete memory update." + ), + ) + + return current_memory + if updated_memory: context.runtime_memory = updated_memory context.runtime_memory_updates = ( @@ -186,6 +356,10 @@ async def summarize_runtime_memory( "[MEMORY] runtime memory updated", ) + await emit_runtime_memory_update( + context + ) + return getattr( context, "runtime_memory", @@ -258,3 +432,35 @@ def schedule_runtime_memory_update( ) return task + + +def schedule_interrupted_runtime_memory_update( + *, + context, +) -> asyncio.Task | None: + + user_message = getattr( + context, + "runtime_turn_user_message", + "", + ) + + assistant_message = ( + build_interrupted_assistant_message( + user_message=user_message, + assistant_message=getattr( + context, + "runtime_turn_assistant_response", + "", + ), + ) + ) + + if not user_message.strip(): + return None + + return schedule_runtime_memory_update( + context=context, + user_message=user_message, + assistant_message=assistant_message, + ) diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py index 9241fc00..60c80bd6 100644 --- a/runtime/runtime_context.py +++ b/runtime/runtime_context.py @@ -53,3 +53,9 @@ class RuntimeContext: background_tasks: set = field( default_factory=set ) + + runtime_turn_user_message: str = "" + + runtime_turn_assistant_response: str = "" + + runtime_turn_interrupted: bool = False diff --git a/runtime/runtime_stream.py b/runtime/runtime_stream.py index 78c14fb9..c5b1e720 100644 --- a/runtime/runtime_stream.py +++ b/runtime/runtime_stream.py @@ -170,6 +170,15 @@ def record_token_usage(self): ), ) + def capture_runtime_turn_response(self): + + if not self.is_brain_context(): + return + + self.context.runtime_turn_assistant_response = ( + self.stream.response + ) + def build_action_log( self, action_event_offset: int, @@ -306,6 +315,8 @@ async def run( ) if not is_valid: + self.capture_runtime_turn_response() + await self.stream.finish( emit=self.emit_to_chat ) @@ -313,6 +324,7 @@ async def run( return None await self.refresh_token_usage() + self.capture_runtime_turn_response() await self.stream.finish( emit=self.emit_to_chat @@ -320,6 +332,7 @@ async def run( await self.refresh_token_usage() self.record_token_usage() + self.capture_runtime_turn_response() log_response = self.stream.response @@ -340,6 +353,9 @@ async def run( except asyncio.CancelledError: + self.context.runtime_turn_interrupted = True + self.capture_runtime_turn_response() + await self.logger.log_runtime( f"{self.runtime_id} stream cancelled." ) @@ -372,6 +388,9 @@ async def run( httpx.RemoteProtocolError, ): + self.context.runtime_turn_interrupted = True + self.capture_runtime_turn_response() + await self.logger.log_system( "Generation aborted." ) diff --git a/static/socket.js b/static/socket.js index f75a81f2..5f433932 100644 --- a/static/socket.js +++ b/static/socket.js @@ -8,6 +8,9 @@ const chatForm = const userInput = document.getElementById("user-input"); +const stopIndicator = + document.getElementById("stop-indicator"); + const sendButton = chatForm.querySelector( 'button[type="submit"]' @@ -46,6 +49,68 @@ function setGenerationState( generationRunning = active; + userInput.readOnly = + active; + + chatForm.setAttribute( + "aria-busy", + active + ? "true" + : "false" + ); + + chatForm.classList.toggle( + "cursor-pointer", + active + ); + + chatForm.classList.toggle( + "border-red-400/80", + active + ); + + chatForm.classList.toggle( + "bg-red-950/35", + active + ); + + chatForm.classList.toggle( + "shadow-[0_0_0_1px_rgba(248,113,113,0.25)]", + active + ); + + userInput.classList.toggle( + "placeholder-red-200/70", + active + ); + + userInput.classList.toggle( + "text-red-100", + active + ); + + userInput.classList.toggle( + "cursor-pointer", + active + ); + + userInput.classList.toggle( + "caret-transparent", + active + ); + + if (stopIndicator) { + stopIndicator.classList.toggle( + "hidden", + !active + ); + + stopIndicator.classList.toggle( + "flex", + active + ); + } + if (!sendButton) { return; } @@ -200,7 +265,26 @@ userInput.addEventListener( ); // -------------------------------------------------- -// BUTTON CLICK +// STOP AREA CLICK +// -------------------------------------------------- + +chatForm.addEventListener( + "click", + function (e) { + + if (!generationRunning) { + return; + } + + e.preventDefault(); + + abortGeneration(); + + } +); + + +// HIDDEN BUTTON CLICK // -------------------------------------------------- if (sendButton) { @@ -247,9 +331,17 @@ ws.onmessage = function (event) { event.data ); - handleTelemetryMessage( - data - ); + if (window.handleTelemetryMessage) { + window.handleTelemetryMessage( + data + ); + } + + if (window.handleRuntimeMemoryMessage) { + window.handleRuntimeMemoryMessage( + data + ); + } // ----------------------------- // LOGS diff --git a/static/telemetry.js b/static/telemetry.js index 5283fb5a..bfd7d29b 100644 --- a/static/telemetry.js +++ b/static/telemetry.js @@ -76,6 +76,16 @@ const contextRuntimePanel = "context-runtime-panel" ); +const runtimeMemoryText = + document.getElementById( + "runtime-memory-text" + ); + +const runtimeMemoryCount = + document.getElementById( + "runtime-memory-count" + ); + function findRuntimeByLabel( runtimes, @@ -444,6 +454,31 @@ function updateChatHeader( } +function renderRuntimeMemory( + memory, + updates +) { + + if (runtimeMemoryText) { + runtimeMemoryText.textContent = + ( + memory + && memory.trim() + ) + ? memory.trim() + : "No runtime memory yet."; + } + + if (runtimeMemoryCount) { + runtimeMemoryCount.textContent = + String( + updates || 0 + ); + } + +} + + function renderContextPanel() { if (isContextTabDisabled( @@ -632,6 +667,23 @@ window.handleTelemetryMessage = function (data) { }; +window.handleRuntimeMemoryMessage = function (data) { + + if ( + !data + || data.type !== "runtime_memory_update" + ) { + return; + } + + renderRuntimeMemory( + data.memory || "", + data.updates || 0 + ); + +}; + + if (window.jinLatestStatus) { window.updateRuntimePanelFromStatus( diff --git a/templates/index.html b/templates/index.html index 2fd39355..0ede7b1b 100644 --- a/templates/index.html +++ b/templates/index.html @@ -141,15 +141,15 @@

+ class="border-t border-slate-500 bg-slate-700 px-5 py-4 shrink-0">
+ class="relative max-w-4xl mx-auto flex items-center bg-slate-600 border border-slate-500 rounded-lg px-4 py-2 transition">
+
+ +
+ +
+ + + [ runtime memory ] + + + + 0 + + +
+ + 
User and JIN just started interacting.
+ +
+ +
+
@@ -344,12 +382,12 @@

"runtimeConfig": runtime_config, } | tojson }} - + - + diff --git a/tests/test_message_memory.py b/tests/test_message_memory.py index 27699e25..d45571a9 100644 --- a/tests/test_message_memory.py +++ b/tests/test_message_memory.py @@ -5,20 +5,30 @@ build_brain_system_prompt, ) from memory.message_memory import ( + build_interrupted_assistant_message, build_runtime_memory_system_prompt, + schedule_interrupted_runtime_memory_update, schedule_runtime_memory_update, summarize_runtime_memory, ) +from settings.config_loader import ( + config, +) class FakeServiceClient: def __init__( self, - response_text: str, + response_text, + finish_reasons=None, ): self.response_text = response_text + self.finish_reasons = list( + finish_reasons + or [] + ) self.calls = [] async def ask( @@ -37,13 +47,33 @@ async def ask( "max_tokens": max_tokens, }) + if isinstance( + self.response_text, + list, + ): + content = self.response_text[ + len( + self.calls + ) + - 1 + ] + else: + content = self.response_text + + choice = { + "message": { + "content": content, + }, + } + + if self.finish_reasons: + choice["finish_reason"] = ( + self.finish_reasons.pop(0) + ) + return { "choices": [ - { - "message": { - "content": self.response_text, - }, - }, + choice, ], } @@ -97,11 +127,71 @@ def test_runtime_memory_prompt_focuses_on_summary_depth(self): "Use deep summarization", prompt, ) + self.assertIn( + "atomic bullet lines", + prompt, + ) + self.assertIn( + "one semantic entity per line", + prompt, + ) + self.assertIn( + "compact semantic label", + prompt, + ) + self.assertIn( + "Keep memory actionable", + prompt, + ) + self.assertIn( + "interrupted response", + prompt, + ) + self.assertIn( + "do not treat it as resolved", + prompt, + ) + self.assertNotIn( + "space exploration costs", + prompt, + ) + self.assertNotIn( + "assistant established", + prompt, + ) self.assertNotIn( "after one completed user/JIN turn", prompt, ) + def test_interrupted_assistant_message_marks_incomplete(self): + + message = build_interrupted_assistant_message( + user_message="Tell me a story.", + assistant_message="Once upon a", + ) + + self.assertIn( + "interrupted by the user", + message, + ) + self.assertIn( + "incomplete", + message, + ) + self.assertIn( + "Do not treat this turn as resolved", + message, + ) + self.assertIn( + "Tell me a story.", + message, + ) + self.assertIn( + "Once upon a", + message, + ) + def test_brain_prompt_includes_runtime_memory(self): context = SimpleNamespace( @@ -140,11 +230,22 @@ async def test_summarizer_updates_runtime_memory(self): clients={ "service": service_client, }, + emitter=SimpleNamespace( + events=[], + emit=None, + ), logger=logger, runtime_memory="User and JIN just started interacting.", runtime_memory_updates=0, ) + async def emit(event): + context.emitter.events.append( + event + ) + + context.emitter.emit = emit + updated_memory = await summarize_runtime_memory( context=context, user_message="Do you remember this?", @@ -167,6 +268,10 @@ async def test_summarizer_updates_runtime_memory(self): "Do you remember this?", service_client.calls[0]["user_prompt"], ) + self.assertIn( + "atomic bullet lines", + service_client.calls[0]["user_prompt"], + ) self.assertEqual( service_client.calls[0]["system_prompt"], build_runtime_memory_system_prompt(), @@ -175,6 +280,93 @@ async def test_summarizer_updates_runtime_memory(self): "[MEMORY] runtime memory updated", logger.service_logs, ) + self.assertEqual( + context.emitter.events, + [ + { + "type": "runtime_memory_update", + "memory": "The user is testing live runtime memory.", + "updates": 1, + }, + ], + ) + + async def test_summarizer_uses_service_max_tokens(self): + + service_client = FakeServiceClient( + "- Active topic: available functions\n" + "- Capabilities listed: answering questions and writing text", + ) + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=FakeLogger(), + runtime_memory="Initial memory.", + runtime_memory_updates=0, + ) + + updated_memory = await summarize_runtime_memory( + context=context, + user_message="What can you do?", + assistant_message="I can answer questions and write text.", + ) + + self.assertEqual( + updated_memory, + ( + "- Active topic: available functions\n" + "- Capabilities listed: answering questions and writing text" + ), + ) + self.assertEqual( + len( + service_client.calls + ), + 1, + ) + self.assertEqual( + service_client.calls[0]["max_tokens"], + config.SERVICE_MAX_TOKENS, + ) + + async def test_summarizer_skips_incomplete_memory(self): + + logger = FakeLogger() + service_client = FakeServiceClient( + "- Active topic: available functions\n" + "- Capabilities listed: answering questions (emails", + ) + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=logger, + runtime_memory="Initial memory.", + runtime_memory_updates=0, + ) + + updated_memory = await summarize_runtime_memory( + context=context, + user_message="What can you do?", + assistant_message="I can answer questions.", + ) + + self.assertEqual( + updated_memory, + "Initial memory.", + ) + self.assertEqual( + context.runtime_memory, + "Initial memory.", + ) + self.assertEqual( + context.runtime_memory_updates, + 0, + ) + self.assertTrue( + logger.errors + ) async def test_scheduled_update_is_background_task(self): @@ -219,6 +411,52 @@ async def test_scheduled_update_is_background_task(self): 0, ) + async def test_interrupted_update_uses_partial_response(self): + + service_client = FakeServiceClient( + "- Active topic: storytelling\n" + "- Interrupted response: user stopped the answer before completion" + ) + context = SimpleNamespace( + clients={ + "service": service_client, + }, + logger=FakeLogger(), + runtime_memory="Initial memory.", + runtime_memory_updates=0, + runtime_turn_user_message="Tell me a story.", + runtime_turn_assistant_response="Once upon a", + ) + + task = schedule_interrupted_runtime_memory_update( + context=context, + ) + + self.assertIsNotNone( + task + ) + + await task + + user_prompt = service_client.calls[0]["user_prompt"] + + self.assertIn( + "interrupted by the user", + user_prompt, + ) + self.assertIn( + "Tell me a story.", + user_prompt, + ) + self.assertIn( + "Once upon a", + user_prompt, + ) + self.assertIn( + "Interrupted response", + context.runtime_memory, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_runtime_stream_tokens.py b/tests/test_runtime_stream_tokens.py index c74ef11d..c6c6babe 100644 --- a/tests/test_runtime_stream_tokens.py +++ b/tests/test_runtime_stream_tokens.py @@ -1,4 +1,5 @@ import unittest +import asyncio from types import SimpleNamespace from runtime.runtime_registry import ( @@ -100,6 +101,16 @@ async def fake_generator(): } +async def fake_cancelled_generator(): + + yield { + "type": "content", + "content": "partial answer", + } + + raise asyncio.CancelledError() + + class RuntimeStreamTokenTests(unittest.IsolatedAsyncioTestCase): async def test_runtime_context_counter_grows_during_stream(self): @@ -192,6 +203,51 @@ async def test_runtime_context_counter_grows_during_stream(self): status=original_state["status"], ) + async def test_cancelled_brain_stream_captures_partial_response(self): + + runtime_id = settings.SERVICE_MODEL_UID + context = SimpleNamespace( + websocket=FakeWebSocket(), + logger=FakeLogger(), + emitter=FakeEmitter(), + runtime_action_events=[], + runtime_usage_events=[], + runtime_turn_assistant_response="", + runtime_turn_interrupted=False, + ) + + stream = RuntimeStream( + context=context, + runtime_id=runtime_id, + role="service", + context_window=( + settings.SERVICE_CONTEXT_WINDOW + ), + log_method=( + context.logger.log_service + ), + context_snapshot={ + "context_role": "brain", + "system_prompt": "system prompt", + "user_prompt": "user payload", + }, + ) + + result = await stream.run( + fake_cancelled_generator() + ) + + self.assertIsNone( + result + ) + self.assertTrue( + context.runtime_turn_interrupted + ) + self.assertEqual( + context.runtime_turn_assistant_response, + "partial answer", + ) + async def test_non_brain_stream_does_not_update_context_counter(self): runtime_id = settings.SERVICE_MODEL_UID diff --git a/websocket.py b/websocket.py index ecf614d7..6e48baba 100644 --- a/websocket.py +++ b/websocket.py @@ -29,6 +29,7 @@ ) from memory.message_memory import ( + schedule_interrupted_runtime_memory_update, schedule_runtime_memory_update, ) @@ -202,6 +203,10 @@ async def process_message( ) ) + context.runtime_turn_user_message = user_text + context.runtime_turn_assistant_response = "" + context.runtime_turn_interrupted = False + state = AgentState( user_input=user_text ) @@ -294,6 +299,8 @@ async def cancel_current_task( if context: + context.runtime_turn_interrupted = True + active_streams = ( getattr( context, @@ -328,6 +335,11 @@ async def cancel_current_task( "Generation cancelled." ) + if context: + schedule_interrupted_runtime_memory_update( + context=context, + ) + # --------------------------------------------------------- # WEBSOCKET CHAT