UN-1031 [FEAT] In Prompt Studio, provide a way to stop a running prompt - #2283
Open
praveen-formido wants to merge 8 commits into
Open
praveen-formido wants to merge 8 commits into
praveen-formido wants to merge 8 commits into
Conversation
A Prompt Studio run could not be stopped. Once Run (or Run All) was clicked the spinner stayed up until a socket event arrived or a 16-minute client timer gave up, while the backend kept spending OCR and LLM tokens to completion. Closing the tab changed nothing. Nothing in the stack could be interrupted from outside: the PG queue exposes only send/read/set_vt/delete with no revoke or dead-letter, and the consumer runs a claimed task eagerly in-process with no time limit. So the stop is cooperative. The backend records an intent in Redis keyed on run_id — the one identifier that spans extract, index and answer_prompt, and that the browser knows before the first response comes back — and the pipeline checks it where stopping is cheap. - unstract-core owns the vocabulary: the key, is_cancelled, and the sentinel every layer matches on. Reads are best effort, so an unreachable Redis reads as "not cancelled" and never kills a live run. - POST prompt-studio/<pk>/cancel/ records the intent. Naming prompts stops just those; naming none stops the whole run including the extract/index stages its prompts share. - The blocking stages unwind the request itself, clearing the document's indexing flag on the way out — left set, it would block every prompt on that document for INDEXING_FLAG_TTL. - The consumer drops a stopped run before starting it and fires its on_error, because the continuations live in the payload and would die silently with a deleted row. - The executor checks between prompts and before each billable stage. A stopped bulk run returns SUCCESS carrying the answers that finished: reporting failure would send the callback down the error path and throw away work the user already paid for. Tokens spent before the stop are still billed. - The UI replaces a running prompt's Run buttons with Stop, adds Stop All beside Run All, and drops queued-but-unsent runs locally. Cloud follow-ups (single-pass and Simple Prompt Studio run ids, and mid-run checkpoints inside the table/agentic executors) are not in this change; those runs still stop only at the pre-run check. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Caz4hivG25nLnnZZD6aULh
The first pass of this ticket stopped the *next* call. A prompt wedged inside a slow LLM or OCR call kept its worker occupied until that call returned on its own, while the UI sat on "Stopping…" and then, after sixteen minutes, misreported it as a timeout. Stop now abandons whatever is running. Two facts set the shape of the fix. A synchronous `litellm.completion` parks a thread on a socket read and nothing outside that thread can interrupt it, so only litellm's async entry point yields a cancellable task whose cancellation closes the request. Embedding and OCR need no async: llama-index calls the embedding leaf once per batch, and LLMWhisperer is a poll loop, so both abort by raising at a boundary we already control. Abort is a predicate, not a dependency. `unstract/sdk1` is shared with the workflow and API-deployment paths and must not learn about Prompt Studio or Redis, so the executor injects a `() -> bool` closure and the SDK only ever calls it. A caller with no Stop button passes none and keeps today's code path byte for byte. - `sdk1.utils.aborting`: `AbortedError` (deliberately not a TimeoutError, or the retry loop would classify a stop as transient), `abort_scope` as a ContextVar for adapters built deep inside plugin code, sliced sleeps, and `run_abortable` over one process-wide event loop — litellm caches async HTTP clients in a module global, so a closed loop poisons every later call; the loop is reset after a fork because the consumer forks its children. - `retry_utils`: check before each attempt and sleep the backoff in slices, so a stop is not swallowed by up to a minute of delay. - `llm.py`: with a predicate, run `acompletion` on that loop and cancel it; `UNSTRACT_LLM_ABORT_INFLIGHT=false` forces the synchronous path back if a provider misbehaves on async. - LLMWhisperer v2 drives its own status polling (`wait_for_completion=False`) rather than blocking inside the vendored client's uninterruptible loop; v1's loop was already ours. - `legacy_executor`: builds the predicate for IDE runs only, memoizing negatives since the SDK polls twice a second, and translates `AbortedError` into the existing `ExecutionCancelled` — so the consumer, callbacks, sentinel, socket event and UI are unchanged. A stop during indexing removes the half-written nodes, or the document would later read as indexed. - Frontend: after thirty seconds the button says the step in progress is still finishing, and a run the user stopped is reported as stopped rather than as a timeout. What this does not promise, and is stated in the README and the QA criteria: closing our end of a request does not stop the provider processing it or reliably avoid being billed; LLMWhisperer has no cancel endpoint at all; and an aborted call returns no response, so its usage is never read and that spend does not reach the usage rows. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014EsA8TZB83nkXcAe6uR7wL
Raised by QA against the stage map: lookup enrichment was missing from it, and the code had two matching holes. Lookup enrichment is a billable LLM call, and the post-processing webhook fires immediately after it, with no checkpoint in front of either. A stop landing there let both run — and a webhook is the one thing in this pipeline that cannot be taken back once another system has received it. There is now a checkpoint before the pair. The bridge's broad `except Exception`, which exists to degrade gracefully on plugin contract drift, also swallowed `AbortedError`. A stop was logged as "lookup failed, continuing without enrichment" and the run carried on to the webhook — the same mistake already fixed in `_handle_index` and `_handle_summarize`. Aborts now pass through; plugin failures still degrade as before. Both tests were verified to fail against the unfixed code. The checkpoint test times its stop to land after the model has answered, so the only checkpoint that can catch it is the new one — a run stopped earlier would never reach this part of the prompt and would prove nothing about it. Not fixed here, because it lives in unstract-cloud: the lookup plugin's own `except Exception` around `enricher.run()` (workers/plugins/lookup_enrichment/src/base.py) catches the abort one layer below this bridge, streams a red "Lookup failed" to the user and persists a lookup error onto the output. The new checkpoint covers the common timing; a stop landing once the lookup call is already in flight still needs that plugin to re-raise. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014EsA8TZB83nkXcAe6uR7wL
All five found by a standardised review of the branch, each verified against the code and pinned by a test confirmed to fail without its fix. A stop could DELETE A COMPLETE INDEX. `is_document_indexed` probes the vector store by embedding a query, and that embedding call is itself abortable. By the time it runs, `doc_id`, `index` and `vector_db` are all set, so the stop handler called `delete_nodes` — destroying vectors a PREVIOUS run wrote and this one never touched. A `wrote_nodes` flag, set immediately before `perform_indexing`, scopes cleanup to an index this run actually began writing. The prompt card's Stop was too weak on a run it was the only prompt of. Naming `prompt_ids` deliberately spares the extraction and indexing a run's prompts SHARE, because the others still need them — but with no sibling left running there is nothing to spare, and the stop sat out the whole of extraction. Manual testing measured the gap: 9.5s for a per-prompt stop against 2.4s for Stop All, at the same moment of the same extraction. It now sends a whole-run cancel when no other prompt of that run is still going. Stopping one prompt of a bulk run disabled its siblings' Stop buttons. `markRunsStopping` marked the whole run, so every prompt in it read as stopping while the backend kept running and billing them — and with one active run, Stop All went too. The user was left with work they could see costing money and no control that would stop it. Stopping is now tracked per prompt. A cancel arriving as a failure RESULT was reported as an error. When a handler converts a stop rather than raising, `ide_prompt_complete` and `ide_index_complete` emitted `status="failed"` carrying the sentinel as error text, so the user got a red toast for a run they stopped themselves. Both now make the same check `ide_prompt_error` already made. The progress message counted prompts that never ran as done. A whole-run stop breaks the loop, leaving unreached prompts out of both lists, so subtracting reported them as answered: stopping at prompt 2 of 10 streamed "Stopped by user after 9 of 10 prompts". It now counts what actually answered. Verified: workers 1542 passed, frontend 613 passed. Each new test was run against the unfixed source first and observed to fail. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TRwXE5CmWoLy3gjSq8Pm4H
Contributor
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Both P1 findings from the Greptile review of PR #2283, fixed together in one commit so the bot re-reviews once rather than per fix. CANCELLATION BYPASSED RUN AUTHORIZATION. A run id is minted in the browser and reaches the endpoint unverified; nothing tied it to anything server-side, and `request_cancel` consumes only the organisation and the run id. The endpoint checked that the caller could reach the tool in the URL, not that the run belonged to it — so a user who learned another run's id, in a project they may only be able to view, could stop its billable work. There was no relational record to check against, so one is now kept beside the cancel intent: `ps:run-tool:{org}:{run_id}` is written at each of the three dispatch sites and read at cancel. A run owned by a different tool is refused with 403, and refused for the whole request — the same all-or-nothing rule the id validation already followed, because a half-applied cancel leaves the caller unable to tell which runs died. An unrecorded owner stays permissive. "Unknown" is not "not ours": Redis may be unreachable, or the run may predate the record. Refusing then would break legitimate stops during a blip, and buys nothing — with Redis down `request_cancel` cannot record the intent either, so the cancel fails anyway and is reported in `failed`. The permission is deliberately left as is. Cancelling requires exactly what running requires, which is the coherent rule: whoever can start the work can stop it. The gap was the missing binding, not the permission level. STOP ALL CROSSED PROJECTS. `activeRuns` is global and outlives navigating between projects, and carried no record of which tool owned each run. A run started in project A, then Stop All pressed in project B, sent A's run ids through B's endpoint — and the backend accepted them. Runs now record their `toolId`, and both Stop All and the per-prompt Stop consider only the project on screen. With the binding above, the backend would now refuse the cross-project ids anyway; this stops them being sent. Verified: workers + core 1591 passed, frontend 614 passed, backend prompt_studio 132 passed (17 pre-existing errors, all DB-bound tests needing live Postgres, unchanged from before this branch). Each new test was run against the unfixed source first and observed to fail. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TRwXE5CmWoLy3gjSq8Pm4H
Greptile's second review found that the ownership record added in 19e026b was itself bypassable, by the same mechanism it was meant to guard. It is right. `remember_run_owner` did an unconditional SET, so the binding was only as trustworthy as its last writer. A caller who knew a live run id could dispatch their own work under it through a tool they control, taking ownership of someone else's run, and then cancel it through that tool — passing the authorization check on the way through. The binding is now immutable: SET NX, first writer wins. A second tool's attempt leaves the original owner in place, so the cancel is refused as it should have been. Dispatch now also refuses a run id owned by another tool, rather than merely declining to rebind it. Without that the foreign run still EXECUTES under the victim's run id — two runs sharing one id, one cancellation key and one stream of socket events. That is no longer a security problem once the binding holds, but it is a collision nobody should have to debug. The fail-open rule is unchanged and deliberate: a signal-store failure returns True and the dispatch proceeds. Refusing to start work because Redis blinked would be worse than leaving one run unverifiable, and it is the same principle the rest of the feature rests on. Verified: workers + core 1596 passed, backend prompt_studio 132 passed (17 pre-existing errors, DB-bound tests needing live Postgres, unchanged). Five new core tests, including the rebinding attack itself and the legitimate same-tool retry that must keep working. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TRwXE5CmWoLy3gjSq8Pm4H
The branch had fallen nine commits behind, which is why pre-commit.ci was failing: it could not produce a merge to run the hooks against, and reported "error during mergeable check" rather than any hook result. Two files conflicted. `internal_views.py` — main's UN-3815 added an org-scoping guard that 404s when the scoped queryset resolves fewer prompts than were asked for. UN-1031 drops prompts with no answer on a stopped run, so a stop never saves a blank over a previous good answer. Both are needed, and the ORDER IS LOAD-BEARING: the guard compares against the ids the caller asked for, and the stop filter deliberately shrinks that set. Filtering first would fire the guard on every stopped run, answering 404 and discarding exactly the partial results the stop exists to preserve. The guard runs first, the filter second, and a comment says why so it is not reordered later. `prompt_studio_helper.py` — main renamed `success` to `status_result` and switched the check to `ExtractionStatusResult.OK`; UN-1031 adds an early raise for the cancel sentinel ahead of that call. Kept both. Verified on the merged tree: workers + core 1738 passed, frontend 614 passed, backend prompt_studio 132 passed. The backend error count rises from 17 to 57 because main brings new DB-bound test files; all of them, main's included, fail with the same "no schema has been selected" configuration error, unrelated to this merge. Main's tests for the org-scoping guard are among those, so that resolution was verified by reading the merged source rather than by running them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TRwXE5CmWoLy3gjSq8Pm4H
praveen-formido
requested review from
hari-kuriakose,
harini-venkataraman and
muhammad-ali-e
September 15, 2026 07:09
Contributor
|
PR #2283 Review — UN-1031: Stop Running Prompts in Prompt Studio Architecture The PR implements cooperative cancellation across the full stack: Redis-backed intent recording → backend checkpoints → consumer pre-run checks → executor stage boundaries → SDK abort scope The "fail-open" design (Redis down = not cancelled) is the correct choice — no healthy run can be killed by a cache outage. Critical Issues
_get_client() and _reset_client() mutate globals (_client_singleton, _client_last_failure) without a lock. In Django's multi-threaded request handling, two threads can race on initialization or reset. Consequence is mild (worst case: a duplicate client or missed reset), but the fix is trivial — add a threading.Lock to match the existing pattern.
"prompt_ids": prompt_ids + cancelled_prompt_ids Python's operator precedence makes this correct (+ binds tighter than if), but the formatting strongly invites future editors to misread or break it. Should use explicit parentheses: "prompt_ids": (prompt_ids + cancelled_prompt_ids) Important Observations
complete_vision previously called litellm.completion(...) directly, bypassing pop_litellm_retry_kwargs. This meant max_retries was passed to litellm AND applied by the retry wrapper = double retries.
prompt_ids + cancelled_prompt_ids could contain duplicates if a race causes a prompt to appear in both lists. The frontend uses includes() so it's not functionally broken, but a Set-based What's Done Well
Verdict Well-architected PR with thoughtful design decisions. The two issues worth fixing before merge are the missing thread lock on the Redis singleton and the misleading ternary formatting — both |
Two points from Harini's review of PR #2283. The cancellation Redis client is built lazily into module globals, and `_get_client` / `_reset_client` mutated them with no lock. Django serves requests on several threads, so two could race to build the client, or interleave a reset with a build and lose the reset. Mild in effect, but the fix is small and matches `aborting._get_loop`, which already guards the same shape with a lock. The common path stays lock-free, since checkpoints call this in tight per-prompt loops and a global read is atomic; only the build and the reset take the lock, and the build re-checks under it. A new test races eight threads through a slow client factory: without the lock it built the client eight times, with it once. `"prompt_ids": prompt_ids + cancelled_prompt_ids if was_cancelled else prompt_ids` was correct only because `+` binds tighter than the conditional, and read as though it might not be. Parenthesised. Not changed: deduplicating that concatenation. The two lists cannot overlap, because `prompt_ids` is filtered a few lines above to exclude every id in `cancelled_prompt_ids`. The `complete_vision` retry behaviour change she noted is documented in the PR description instead, as it needs no code change. Verified: workers + core 1739 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Caz4hivG25nLnnZZD6aULh
Contributor
Frontend Lint Report (Biome)✅ All checks passed! No linting or formatting issues found. |
Contributor
Author
|
@harini-venkataraman thanks for the review. I've addressed it in 1. Thread-safety of the Redis client singleton: fixed. 2. Misleading ternary formatting: fixed. 3. 4. Deduplicating All workers and core tests pass: 1739. Separately, |
|
Contributor
Unstract test resultsPer-group results
Critical paths
|
harini-venkataraman
approved these changes
Sep 16, 2026
harini-venkataraman
left a comment
harini-venkataraman
left a comment
Contributor
There was a problem hiding this comment.
LGTM, tests are failing, please check and merge.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes UN-1031. Scope is Prompt Studio only — the same capability for API deployments is planned separately in UN-4102.
Prompt Studio gave no way to stop a prompt once it was running. A user who picked the wrong model or document, or simply changed their mind, had to sit and watch it — and pay for it. This adds a Stop on each prompt card and a Stop All in the toolbar.
How it works
Nothing can interrupt a running task from outside: the PG queue has no revoke, Celery's
revoke(terminate=True)is inert here (childrenSIG_IGNSIGTERM, the worker uses thethreadspool), and a synchronouslitellm.completionparks its thread on a socket read. So the design inverts control — the canceller writes a fact, and the run reads it:The abort predicate is a
Callable[[], bool]closure the executor injects, sounstract/sdk1never learns about Prompt Studio or Redis. Workflow and API-deployment runs pass no predicate and keep the previous code path byte for byte.Key mechanics: one process-wide event loop (litellm caches async clients keyed to the loop they were built on), reset via
os.register_at_forkbecause the consumer forks children; sliced sleeps so a stop isn't swallowed by retry backoff;AbortedErrordeliberately not aTimeoutErrorsubclass, or the retry predicate would retry the call we just walked away from; negative predicate results memoised for 1s, since a 0.5s poll across a 900s call is ~1800 Redis round trips.UNSTRACT_LLM_ABORT_INFLIGHT=falsereverts to the synchronous path if a provider misbehaves on async.What it does not promise
Closing our end of a request does not stop the provider processing it or reliably avoid the bill. LLMWhisperer has no cancel endpoint, so an abandoned extraction runs to completion on their side — observed in testing, completing 2.5s after our stop was recorded. And an aborted call returns no response, so its token counts never reach the usage rows. All three are documented in the module docstring, the pg_queue README and the QA criteria.
Which button stops what
Stop Allcancels the whole run including the extraction and indexing its prompts share. A per-promptStopdeliberately spares those shared stages, because the run's other prompts still need them. Measured on a local stack: Stop All during extraction ended the run in 2.4s; a per-prompt stop at the same moment took 9.5s. Both correct — worth knowing before testing stages 3–5.Review remediation
A standardised review found five defects, all fixed in
7eecb4142, each pinned by a test confirmed to fail without its fix:is_document_indexedembeds a query to probe the store, and that call is abortable — by thendoc_idnames a document a previous run indexed, so cleanup destroyed vectors this run never wrote.Behaviour change outside the stop path
LLM.complete_visionnow goes through_invoke_completion, the same path ascomplete, so a stop can abandon a vision call too. As a side effect its retry behaviour changes: onmainit calledlitellm.completiondirectly, with the adapter'smax_retriespassed through and applied by litellm itself; it now hasmax_retriespopped and applied by the SDK'scall_with_retry, which uses our retryable-error predicate and backoff, the same policycompletealready uses. The attempt budget is still the adapter'smax_retries, but which errors are retried and how long it waits between them now matchcomplete— and for providers where litellm did not honourmax_retriesitself, vision calls now actually get retried. (Raised by @harini-venkataraman in review.)Verification
test_aborting,test_retry_abort,test_llm_abort,test_whisperer_v2_abort, plus executor/consumer/callback cancellation testsQA acceptance criteria: https://claude.ai/code/artifact/fcd0c433-4d59-4b41-a03b-9c3583ce1207
Known gap, needs a cloud-side change
The lookup plugin lives in
unstract-cloud. Its own broad handler inworkers/plugins/lookup_enrichment/src/base.pycatches the abort below this repo's bridge fix, so a stop landing while a lookup call is already in flight still shows "Lookup failed" and writes that error onto the saved output. The new checkpoint covers the common timing; the passthrough there is a separate two-line PR.Before merge
One clean run per configured LLM provider — Bedrock, Claude and VertexAI at minimum. Prompt Studio now reaches providers through litellm's async entry point, so an authentication or transport difference would surface there and nowhere else. This is the highest-risk item in the change.
🤖 Generated with Claude Code
https://claude.ai/code/session_01TRwXE5CmWoLy3gjSq8Pm4H