Skip to content

fix(services): detach the turn's status callback when the turn ends #264

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dovvnloading merged 1 commit into from
Sep 7, 2026
Merged

fix(services): detach the turn's status callback when the turn ends #264

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
260 changes: 138 additions & 122 deletions backend/cortex_backend/services/generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,13 @@ class GenerationEngine(Protocol):
last_code_proposal: CodeExecutionProposal | None
last_code_rejection: CodeProposalRejection | None

def set_status_callback(self, callback: Callable[[str], None]) -> None:
"""Receive startup progress while ``generate`` blocks."""
def set_status_callback(self, callback: Callable[[str], None] | None) -> None:
"""Receive startup progress while ``generate`` blocks.

``None`` detaches it. The chat client behind a real engine is
process-wide, so a callback left installed keeps the turn that made it
-- and everything its closure holds -- alive past that turn's end.
"""

def fit_memories_to_context(
self,
Expand Down Expand Up @@ -255,137 +260,148 @@ def generate(
engine.set_status_callback(
lambda message: self._publish(sink, snapshot, "loading_model", message)
)
try:

self._check_cancelled(cancellation_event)
loaded_history = (
history_messages
if history_messages is not None
else self._history_loader(snapshot.thread_id)
)
working_history = [dict(message) for message in loaded_history]
if working_history and working_history[-1].get("role") == "user":
working_history.pop()

# Reserve room for attachments *before* history claims the whole
# budget: fit them first against a placeholder (history is not known
# yet), giving an attached document priority over old chat turns,
# then let history size itself around that reservation below. The
# attachments passed to engine.generate() further down are re-fit
# against the real, now-correctly-sized chat_history -- this pass
# only determines how much room history should leave.
reserved_attachments: Sequence[GenerationAttachment] = ()
if snapshot.attachments:
reserved_attachments = engine.fit_attachments_to_context(
snapshot.attachments,
query=snapshot.user_input,
chat_history="No history available.",
permanent_memories=permanent_memories,
memories_enabled=snapshot.memories_enabled,
user_system_instructions=snapshot.user_system_instructions,
num_ctx=num_ctx,
code_execution_eligible=snapshot.code_execution_eligible,
bypass_system_prompt=snapshot.bypass_system_prompt,
host_observations=snapshot.host_observations,
self._check_cancelled(cancellation_event)
loaded_history = (
history_messages
if history_messages is not None
else self._history_loader(snapshot.thread_id)
)
working_history = [dict(message) for message in loaded_history]
if working_history and working_history[-1].get("role") == "user":
working_history.pop()

# Reserve room for attachments *before* history claims the whole
# budget: fit them first against a placeholder (history is not known
# yet), giving an attached document priority over old chat turns,
# then let history size itself around that reservation below. The
# attachments passed to engine.generate() further down are re-fit
# against the real, now-correctly-sized chat_history -- this pass
# only determines how much room history should leave.
reserved_attachments: Sequence[GenerationAttachment] = ()
if snapshot.attachments:
reserved_attachments = engine.fit_attachments_to_context(
snapshot.attachments,
query=snapshot.user_input,
chat_history="No history available.",
permanent_memories=permanent_memories,
memories_enabled=snapshot.memories_enabled,
user_system_instructions=snapshot.user_system_instructions,
num_ctx=num_ctx,
code_execution_eligible=snapshot.code_execution_eligible,
bypass_system_prompt=snapshot.bypass_system_prompt,
host_observations=snapshot.host_observations,
)

history_kwargs: dict[str, Any] = {
"query": snapshot.user_input,
"permanent_memories": permanent_memories,
"memories_enabled": snapshot.memories_enabled,
"user_system_instructions": snapshot.user_system_instructions,
"num_ctx": num_ctx,
"code_execution_eligible": snapshot.code_execution_eligible,
"bypass_system_prompt": snapshot.bypass_system_prompt,
"host_observations": snapshot.host_observations,
}
if reserved_attachments:
history_kwargs["attachments"] = reserved_attachments
# Preferred shape: the same retained exchanges as real user/assistant
# turns, which is what a chat-tuned model's template expects and what
# lets a local runtime reuse its cache across turns. One call returns
# both renderings because choosing which exchanges fit is the expensive
# part and must not be done twice. Engines that do not offer it (the
# narrower fakes in the test suite, and any older adapter) keep the
# flattened transcript.
chat_history, structured_history = engine.fit_history(
working_history, **history_kwargs
)

self._check_cancelled(cancellation_event)
generate_kwargs: dict[str, Any] = {
"query": snapshot.user_input,
"chat_history": chat_history,
"permanent_memories": permanent_memories,
"memories_enabled": snapshot.memories_enabled,
"user_system_instructions": snapshot.user_system_instructions,
"options": dict(snapshot.model_options),
"host_observations": snapshot.host_observations,
}
# Keep the legacy headless engine protocol compatible for callers that
# do not use attachments or cancellation; real engines receive the
# resolved payload.
if snapshot.attachments:
generate_kwargs["attachments"] = snapshot.attachments
if cancellation_event is not None:
generate_kwargs["cancellation_event"] = cancellation_event
generate_kwargs["history_messages"] = structured_history
response, thoughts, memory_command, stats = engine.generate(
**generate_kwargs,
)
if not isinstance(memory_command, MemoryCommand):
raise ModelOperationError(
"Generation returned an invalid memory command.",
operation="generation",
history_kwargs: dict[str, Any] = {
"query": snapshot.user_input,
"permanent_memories": permanent_memories,
"memories_enabled": snapshot.memories_enabled,
"user_system_instructions": snapshot.user_system_instructions,
"num_ctx": num_ctx,
"code_execution_eligible": snapshot.code_execution_eligible,
"bypass_system_prompt": snapshot.bypass_system_prompt,
"host_observations": snapshot.host_observations,
}
if reserved_attachments:
history_kwargs["attachments"] = reserved_attachments
# Preferred shape: the same retained exchanges as real user/assistant
# turns, which is what a chat-tuned model's template expects and what
# lets a local runtime reuse its cache across turns. One call returns
# both renderings because choosing which exchanges fit is the expensive
# part and must not be done twice. Engines that do not offer it (the
# narrower fakes in the test suite, and any older adapter) keep the
# flattened transcript.
chat_history, structured_history = engine.fit_history(
working_history, **history_kwargs
)
if not snapshot.memories_enabled:
memory_command = MemoryCommand()

proposal = engine.last_code_proposal
if not snapshot.code_execution_eligible or not isinstance(
proposal, CodeExecutionProposal
):
proposal = None
rejection = engine.last_code_rejection
if not isinstance(rejection, CodeProposalRejection) or proposal is not None:
rejection = None

if snapshot.translation_enabled:

self._check_cancelled(cancellation_event)
self._publish(
sink,
snapshot,
"translation",
f"Translating to {snapshot.target_language}...",
)
translation_result = _call_with_optional_kwargs(
engine.translate_text,
response,
snapshot.target_language,
options=dict(snapshot.model_options),
cancellation_event=cancellation_event,
generate_kwargs: dict[str, Any] = {
"query": snapshot.user_input,
"chat_history": chat_history,
"permanent_memories": permanent_memories,
"memories_enabled": snapshot.memories_enabled,
"user_system_instructions": snapshot.user_system_instructions,
"options": dict(snapshot.model_options),
"host_observations": snapshot.host_observations,
}
# Keep the legacy headless engine protocol compatible for callers that
# do not use attachments or cancellation; real engines receive the
# resolved payload.
if snapshot.attachments:
generate_kwargs["attachments"] = snapshot.attachments
if cancellation_event is not None:
generate_kwargs["cancellation_event"] = cancellation_event
generate_kwargs["history_messages"] = structured_history
response, thoughts, memory_command, stats = engine.generate(
**generate_kwargs,
)
if not isinstance(translation_result, TranslationResult):
if not isinstance(memory_command, MemoryCommand):
raise ModelOperationError(
"Translation returned an invalid result.",
operation="translation",
"Generation returned an invalid memory command.",
operation="generation",
)
if not translation_result.success:
raise ModelOperationError(
translation_result.error or "Translation failed. Please try again.",
operation="translation",
if not snapshot.memories_enabled:
memory_command = MemoryCommand()

proposal = engine.last_code_proposal
if not snapshot.code_execution_eligible or not isinstance(
proposal, CodeExecutionProposal
):
proposal = None
rejection = engine.last_code_rejection
if not isinstance(rejection, CodeProposalRejection) or proposal is not None:
rejection = None

if snapshot.translation_enabled:
self._check_cancelled(cancellation_event)
self._publish(
sink,
snapshot,
"translation",
f"Translating to {snapshot.target_language}...",
)
translation_result = _call_with_optional_kwargs(
engine.translate_text,
response,
snapshot.target_language,
options=dict(snapshot.model_options),
cancellation_event=cancellation_event,
)
response = translation_result.text or ""
if not isinstance(translation_result, TranslationResult):
raise ModelOperationError(
"Translation returned an invalid result.",
operation="translation",
)
if not translation_result.success:
raise ModelOperationError(
translation_result.error or "Translation failed. Please try again.",
operation="translation",
)
response = translation_result.text or ""

self._check_cancelled(cancellation_event)
self._check_cancelled(cancellation_event)

return GenerationServiceResult(
response=response,
thoughts=thoughts,
memory_command=memory_command,
code_execution_proposal=proposal,
code_execution_rejection=rejection,
stats=stats,
)
return GenerationServiceResult(
response=response,
thoughts=thoughts,
memory_command=memory_command,
code_execution_proposal=proposal,
code_execution_rejection=rejection,
stats=stats,
)
finally:
# The chat client is process-wide while the engine is built per
# turn, so this callback outlived the turn that installed it. It
# closes over the snapshot -- attachments included -- which stayed
# referenced for the life of the process, and any later status
# message reached a finished turn: generate_chat_title builds a
# fresh engine and installs no callback of its own, so a model
# load during titling published "loading_model" against the job
# that had already completed.
engine.set_status_callback(None)

def generate_chat_title(
self,
Expand Down
2 changes: 1 addition & 1 deletion backend/cortex_backend/testing/fake_ollama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ def __init__(self, state: FakeOllamaState | None = None):
self.last_code_proposal: CodeExecutionProposal | None = None
self.last_code_rejection: CodeProposalRejection | None = None

def set_status_callback(self, callback: Callable[[str], None]) -> None:
def set_status_callback(self, callback: Callable[[str], None] | None) -> None:
self._status_callback = callback

def fit_memories_to_context(
Expand Down
53 changes: 53 additions & 0 deletions tests/test_services.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,59 @@ def chat(self, *, model, messages, options, cancellation_event=None):
self.assertEqual(seen["model"], "translategemma:4b")
self.assertIs(seen["cancellation_event"], event)

def test_the_status_callback_does_not_outlive_its_turn(self):
"""A per-turn callback must not stay on the process-wide chat client.

`SynthesisAgent.set_status_callback` forwards to the chat client, which
is built once for the process while the engine is built per turn. The
callback closes over the snapshot -- attachments and all -- so leaving
it installed kept that turn alive for the life of the process, and any
later status message reached a finished turn: `generate_chat_title`
builds a fresh engine and installs no callback of its own, so a model
load during titling published "loading_model" against a job that had
already completed.
"""
installed: list[object] = []

class _SharedClientEngine(_FakeEngine):
def set_status_callback(self, callback) -> None:
installed.append(callback)

service = GenerationService(
history_loader=lambda thread_id: [],
memory_loader=lambda: [],
engine_factory=lambda snapshot: _SharedClientEngine(),
)

service.generate(_snapshot(memories_enabled=False, translation_enabled=False))

self.assertTrue(installed, "a callback was never installed at all")
self.assertIsNone(installed[-1], "the turn's callback was left on the shared client")

def test_the_status_callback_is_detached_even_when_the_turn_fails(self):
"""The failure path is the one that matters most: it leaves a
half-finished turn behind, and that is exactly when a stale callback
would misattribute the next runtime message."""
installed: list[object] = []

class _FailingEngine(_FakeEngine):
def set_status_callback(self, callback) -> None:
installed.append(callback)

def generate(self, **kwargs):
raise ModelOperationError("the model fell over", operation="generation")

service = GenerationService(
history_loader=lambda thread_id: [],
memory_loader=lambda: [],
engine_factory=lambda snapshot: _FailingEngine(),
)

with self.assertRaises(ModelOperationError):
service.generate(_snapshot(memories_enabled=False, translation_enabled=False))

self.assertIsNone(installed[-1], "a failed turn left its callback installed")

def test_generation_is_headless_and_emits_owned_typed_progress(self):
engine = _FakeEngine()
recorder = _ProgressRecorder()
Expand Down