Skip to content

Cancel AI assistant in-flight run - #4563

Draft
nelsonkopliku wants to merge 16 commits into
mainfrom
cancel-run-on-new-chat
Draft

Cancel AI assistant in-flight run#4563
nelsonkopliku wants to merge 16 commits into
mainfrom
cancel-run-on-new-chat

Conversation

@nelsonkopliku

@nelsonkopliku nelsonkopliku commented Jul 30, 2026

Copy link
Copy Markdown
Member

Clicking "New chat" mid-stream swapped threadID and cleared the messages, but nothing settled the run: AgUiThreadRuntimeCore kept isRunningFlag set and WebSocketAIAgent kept the old subscriber alive until the server sent RUN_FINISHED. The new thread therefore rendered an optimistic empty assistant message with a spinner and hid the send button, so the user could not prompt until the abandoned answer finished.

Frontend:

  • AssistantChatProvider clears the thread with import({messages: []}) and then cancels a run that is still in flight. import() rather than reset() because the external-store runtime's reset() leaves its own message repository untouched; clear-before-cancel because cancelRun() re-applies a pre-cancel snapshot a macrotask later.
  • The runtime gets an onCancel that routes to the transport, so the abort signal fires before the subscriber settles and the abandoned run stays silent instead of surfacing as a failure.
  • WebSocketAIAgent grows cancelActiveRun() (pushes cancel_run, completes the subscriber) and ignores RUN_STARTED/FINISHED/ERROR whose runId belongs to a run that is no longer active.

Backend:

  • AIAssistantChannel handles cancel_run: cancels the agent's current run and clears :loading. That flag is channel-scoped, so without clearing it the first prompt of the new chat would hit the double-send guard and get dropped with no reply at all.
  • Trento.AI.Agent.cancel/1 delegates to Sagents.AgentServer.cancel/1 rather than stop/1: stopping goes through AgentServer.terminate/2, which waits up to 25s for a running task. A dead agent name would exit the caller, so the call is wrapped in catch :exit.

Description

Fixes #

How was this tested?

Documentation changes

Yes / No

Additional information

@nelsonkopliku nelsonkopliku self-assigned this Jul 30, 2026
@nelsonkopliku nelsonkopliku added this to the 3.2.0 milestone Jul 30, 2026
@nelsonkopliku nelsonkopliku added the enhancement New feature or request label Jul 30, 2026
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch from dac8d78 to 66487ac Compare August 3, 2026 07:40
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch from 66487ac to 9860d38 Compare August 3, 2026 08:07
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch 2 times, most recently from 1289ae8 to 95c558d Compare August 3, 2026 10:43
Base automatically changed from improve-ai-assistant-tests to main August 3, 2026 16:27
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch 2 times, most recently from 5af280b to b82f21c Compare August 5, 2026 15:09
@nelsonkopliku nelsonkopliku mentioned this pull request Aug 5, 2026
Merged
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch 2 times, most recently from 50322f0 to 61e3c4c Compare August 7, 2026 14:20
@nelsonkopliku nelsonkopliku changed the title Cancel the in-flight run when the user starts a new chat Cancel the in-flight run Aug 7, 2026
@nelsonkopliku nelsonkopliku changed the title Cancel the in-flight run Cancel AI assistant in-flight run Aug 7, 2026
nelsonkopliku and others added 9 commits August 10, 2026 09:42
@nelsonkopliku
Cancel in-flight run on New Chat
69441e8
@nelsonkopliku
Split cancel_run from abandon_thread on the assistant channel
17b2ced 
cancel_run now cancels the in-flight run and leaves the thread's agent
standing, so the same thread can be prompted again. The old behaviour —
cancel and terminate — moves to a new abandon_thread event, which is what
"New chat" sends.
@nelsonkopliku @claude
Restore the "New chat" lock while a run is in flight
d937ef7 
Stop in the composer becomes the only way to end a run; "New chat" is
reachable only when the thread is idle. This restores #4581, which this
branch had reverted.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku
Send abandon_thread on "New chat", keep cancel_run for Stop
1120805 
cancelActiveRun now no-ops when nothing is in flight: it is Stop's path
only, and Stop is unreachable without a run. The provider gets its own
abandonThread push for the thread it is walking away from.
@nelsonkopliku
Settle the local run when abandonThread terminates the server agent
bb24dcc 
New chat's threadID has a second writer: a cross-tab
ai_configuration_created can mint one while the launcher is closed, with
a run genuinely still streaming. The server tears the agent down outright
for abandon_thread, so no RUN_FINISHED/RUN_ERROR ever comes back — the
client has to settle its own run, the same contract cancelActiveRun()
already honours. Without it, isRunning wedges true forever on that path.
@nelsonkopliku @claude
Replace Send with a Stop button while a run is in flight
499160e 
The Stop button goes through ComposerPrimitive.Cancel, which reaches the
transport via the runtime's onCancel. It is never disabled: a run can
outlive the connection or the AI configuration, and neither should strand
the user mid-answer.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku @claude
Mark a stopped answer with a subtle notice
90378f6 
The AG-UI run aggregator stamps a cancelled run's message as
incomplete/cancelled, which is the only signal distinguishing "the user
hit Stop" from "the answer ended". reason: 'error' is left to
MessageError.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku
Drive composer Stop and StoppedNotice from our own state, not the run…
2b3eb41 
…time

ComposerPrimitive.Cancel (-> runtime.thread.cancelRun()) and a message-status-
driven StoppedNotice both sit on confirmed @assistant-ui library defects with
no supported extension point: cancelRun() resyncs a stale message snapshot
that wedges thread.isRunning permanently, and react-ag-ui's subscriber
dispatches a synthesized RUN_FINISHED on any settle (including cancel) that
clobbers incomplete/cancelled before anything can read it.

Add StoppedRunProvider, which owns stopRun/isStopped from our own state and
threads them down as props the same way isRunning already is. The composer's
Stop button becomes a plain Button calling onStop directly, and StoppedNotice
now gates on isStopped + s.message.isLast instead of message status.

This replaces the design of tasks 4 and 5 (83fc181, 8830bc5) without
reverting them.
@nelsonkopliku @claude
Pin the composer Stop action end to end
586d603 
Three integration tests over the real agent, runtime and rendering path,
covering what Stop has to do and what it must not do:

- stopping mid-stream pushes a single cancel_run carrying the run's pinned
  ids, keeps the prompt and the partial answer on screen, marks the answer
  as stopped, and hands the composer back empty
- stopping before the first token behaves identically, with no server event
  involved at all
- the thread outlives the stop: the next prompt is a new run on the same
  thread_id

All three fail against the pre-rework code, where the runtime's own cancel
path left the composer wedged on Stop and the marker unreachable.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
nelsonkopliku and others added 7 commits August 10, 2026 09:43
@nelsonkopliku
Gate the stopped marker on render, not on StoppedRunProvider's effect
e4c717b 
isStopped clears from a passive effect, which cannot run before the
next run's placeholder commits. That leaves one committed render where
the new, empty assistant answer is isLast, isStopped is still true from
the previous run, and StoppedNotice shows "Response stopped." on it.
Pass isStopped && !isRunning instead so the render itself withholds the
marker while a run is in flight.
@nelsonkopliku @claude
Make the stopped marker per-message, per the approved spec
6df4c44 
Three review findings, one commit:

- The shipped "Response stopped." marker was a single thread-scoped
  boolean, cleared the instant the next run started, and gated on
  message.isLast — so it vanished from the cut-off answer as soon as the
  user sent another prompt. The approved spec requires it to persist on
  that message for the rest of the conversation. Replace the boolean with
  a Set of stopped message ids keyed off the thread's own messages, and
  delete the machinery that only existed to compensate for it: the isLast
  guard, the clearing effect, the isStopped prop chain through
  MessageBubble/AssistantThread/AIAssistant, and StoppedNotice's isLast
  read.

- Pin the closed-launcher-mid-run path end to end. abandonThread()
  settling the local run (task 3) was previously guarded only by a
  transport unit test asserting subscriber.complete() was called. Add an
  integration test that opens the assistant, starts a run, closes the
  launcher, lets a cross-tab ai_configuration_created abandon the thread,
  and reopens — asserting the composer recovers to "Send message" and
  "New chat" is enabled. Confirmed it fails (composer stuck on
  "Stop generating") by temporarily dropping the _settleActiveRun() call
  from abandonThread().

- cancelActiveRun() now reports whether it actually settled a run
  (false when nothing was in flight), and StoppedRunProvider.stopRun
  gates marking a message on that — a click landing after RUN_FINISHED
  but before the composer swaps Stop→Send can no longer label a
  completed answer as stopped. Fixed a stale comment left over from the
  runtime-onCancel design this branch already moved away from.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku @claude
test: pin the stop verdict and per-message marking in the unit suites
1b417b4 
The final review left two test-side notes after the marker went
per-message:

- AssistantChatProvider.test.jsx stubbed useAuiState to a constant, so
  StoppedRunProvider's marking branch was dead in that suite and the
  mock's comment still described the deleted isRunning read. The stub is
  now selector-driven and StopTrigger reports whether the stop marked the
  thread's last answer, which pins the transport verdict travelling back
  through the provider's stop handler.
- MessageBubble.test.jsx had a test named for a mid-run guarantee that
  rendered without a provider, so it only exercised the default context.
  It now stops a later message and asserts this one stays unmarked, which
  is the per-message guarantee its name claims. The dead isLast in the
  mock state is gone.

Mutation-proved: dropping the return from handleStop fails the first new
test; making isMessageStopped ignore its argument fails the second.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku @claude
fix: show the progress indicator only on the answer in flight
a9d1ced 
isRunning describes the thread, not a message, so every assistant bubble
receives the same value. AgentProgressIndicator told the live placeholder
apart from the rest of the history by checking that the message carried no
text yet — which held only because a run that reached RUN_FINISHED always
left text behind.

Stop broke that assumption. A run stopped before its first token leaves a
permanently text-less assistant message in the conversation, so the next
prompt turned the spinner back on under the answer tagged 'Response
stopped.' as well as under the new one.

Gate on the message being the thread's last, which is the same identity
the library computes for isLast and the same one StoppedRunProvider uses
to decide which answer a stop marks.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@nelsonkopliku
Fix Liz's z-index
746d342
@nelsonkopliku
update ai assistant ui libs
b303a1c
@nelsonkopliku
refactor: drive the composer Stop through ComposerPrimitive.Cancel
e4870ac 
The Stop button is now assistant-ui's own cancel primitive rather than a
plain button wired to the transport, and the "Response stopped." marker
reads the message's own `incomplete/cancelled` status rather than a Set of
message positions we maintain ourselves. `StoppedRunProvider` and the
`onStop` prop threaded through AIAssistant → AssistantThread →
PromptComposer are gone with it.

What makes this possible is the seam `AbstractAgent` leaves for transports:
`AgUiThreadRuntimeCore.cancel()` calls `agent.abortRun()` before aborting
its own controller, and the base implementation is an empty no-op only
`HttpAgent` overrides. `WebSocketAIAgent.abortRun()` supplies the websocket
half — it pushes `cancel_run` with the run's pinned ids and errors the run
Observable with an `AbortError`, the shape the ag-ui adapter reads as
RUN_CANCELLED.

That error is deferred twice, and both hops are load-bearing: a microtask
to land after the runtime's own abort (an error raised before it is treated
as a genuine failure and re-thrown as an unhandled rejection) and past the
synchronous remainder of the external store's `cancelRun()`, then a timer
that lands behind the message resync `cancelRun()` schedules, whose
pre-cancel snapshot would otherwise overwrite the cancelled status.

A run stopped before its first token now unwinds the turn: the empty
exchange is dropped and the prompt goes back into the composer. That is the
library's behaviour for a cancel with nothing streamed yet, and it is a
deliberate change of position from the original design — nothing is lost,
and there is no blank bubble left to mark.
@nelsonkopliku
nelsonkopliku force-pushed the cancel-run-on-new-chat branch from ded7dd9 to e4870ac Compare August 10, 2026 07:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@nelsonkopliku nelsonkopliku

Labels

enhancement New feature or request

Milestone

3.2.0

Development

Successfully merging this pull request may close these issues.

1 participant

@nelsonkopliku