unsloth/studio/backend/storage/studio_db.py
Daniel Han 18b97f8b72
Studio: keep and search the turns rolling context evicts (#9074)
* Studio: add rolling context windows for local GGUF chat

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: return rolling context metadata for non-stream chats

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep original context when rolling fit fails

* Studio: keep instruction groups independently protected

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: preserve rolling context metadata across retries

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: count sanitized rolling context prompts

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: refit rolling context after respawn

* Studio: scope middle truncation to passthrough

* Studio: refit tool prompts after respawn

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: retain later choice truncation metadata

* Studio: report clipping-only context truncation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: expose the turns a rolling fit evicts

Rolling context eviction currently drops turns with no way for a caller to
learn which ones went. Extract the turn grouping so the eviction unit is
reusable, add an identity-based diff of what a fit removed, and let a caller
reserve room for content it intends to add back after fitting.

The reserve deliberately does not affect whether trimming happens at all, only
how far it goes once it is already required, so a conversation that fits today
is still returned untouched.

* Studio: archive evicted turns into a per-thread rag scope

An evicted turn is currently gone for the rest of the session, so the model
will state the conversation began wherever its visible context begins. Keep the
turns the rolling window drops in a searchable scope built on the existing
store, chunker, embedder and hybrid retrieval.

The archive is cumulative: every compaction adds to it and nothing is cleared,
so a later compaction can still find what an earlier one evicted. It lives in
its own scope rather than the thread's document scope, because thread documents
are injected in full on every request and would re-inject the whole history.

Idempotent by content hash, since the same turns are evicted again on every
later request. Every entry point degrades to a no-op rather than raising.

* Studio: recall archived turns on the turn that evicted them

Given only a search tool, a model decides for itself whether to look, and
mostly does not. Measured on MRCR v2, a 35B declined on 56% of rows, scoring
0.099 when it skipped against 0.461 when it searched. Forcing one retrieval on
the compaction turn took tool-only 0.258 to 0.604, and the model then called
the tool on 0% of rows, so the common path costs nothing extra.

Recall fires at most once per request, since the tool loop refits on every
iteration. The tool loop renders it as an ordinary tool exchange through the
builder shared with document auto-inject; the plain path prefixes the latest
user message instead, because it sends no tools array and a tool role without
one breaks strict chat templates.

Only scalar counts join the context_truncated event, so no message content
reaches the wire.

* Studio: let the model search a compacted conversation

Forced recall answers the turn that evicted, but a later turn can refer to
something the forced pass had no reason to fetch. Add search_conversation so
the model can go looking, scoped to the thread's own archive and sharing the
admission slot with document search so the two cannot race for the embedder.

The tool is offered only once a thread has actually had turns archived, so an
ordinary short chat never pays for the schema, and it is classified read-only
so auto mode does not prompt on every call. A matching system-prompt note tells
the model the session was compacted, since otherwise it assumes the
conversation began where its visible context begins.

Deleting a thread now drops its archive rather than leaking a scope per chat.

* Studio: read thread_id defensively when selecting tools

_select_request_tools also serves the token-count request model, which has no
thread_id field, so the archive gate raised there.

* Studio: retrieve archived turns lexically first

Recalling your own conversation is mostly an exact-match problem: a name, a
number, an identifier someone pasted twenty turns ago. Those live or die on
rare-token matching, and hybrid fusion was losing them.

Measured on a 30-turn walkthrough of a 230k-character document at a 16k window,
where every turn shared the same wrapper text. The chunk holding the needle
ranked 3rd lexically at any k, was never returned by dense retrieval at all,
and RRF pushed it to 16th because it had 30 useless dense hits to fuse with.
End to end the model answered with the exact code once lexical leads, and could
not answer at all before.

Dense still fills whatever the lexical pass leaves, for paraphrased recall.

* Studio: keep recall on the branch the user is actually on

Editing an earlier message rewinds a thread and continues down a new branch,
but the archive is append-only and still holds everything the abandoned
continuation produced. Verified against a live build: after rewinding past a
turn, querying its distinctive text still returned it, so the model could be
handed a turn that on this branch never happened.

Recall now drops archived turns that are absent from the thread's saved
transcript. Threads with no saved transcript are left unfiltered, since an API
caller may pass a thread_id without persisting messages and an empty transcript
is absence of evidence rather than evidence the turns are gone.

Containment on a normalised prefix rather than a digest, because the archived
copy is rendered from the inference projection and the saved copy comes back
through the message store.

* Studio: show a persistent compaction notice on the turn that compacted

The only signal that a long chat had been compacted was a toast, which vanishes after
a few seconds and does not survive a reload. A user who scrolls back later has no way
to find out why the model seemed to forget the start of the conversation.

The notice renders from metadata.custom.contextTruncation inside the assistant
message's own container, so it is not part of the conversation sent to the model, is
not editable, and is not exported as content, but it stays attached to the turn it
describes. It reports how many messages were dropped and, when the conversation
archive is on, that they are still searchable and how many passages were recalled.

* Studio: show the compaction notice once, on the turn it started

A thread that has outgrown its window compacts on every turn from then on, not
just the first, so a notice per compacted turn was a notice on every reply for the
rest of the conversation. The user needs telling once.

The notice is now gated on being the first compacted assistant turn in the thread,
found by walking the thread rather than assuming the compacted turns are contiguous
or that this is the last one, so a rollback that removes the turn it was on moves it
to whichever turn now compacts first. The wording follows: it describes the state
the conversation is in from here on, and carries the counts from the turn it began
on in parentheses.

* Studio: make compaction an occasional event instead of every turn

The fit was stateless. The client re-sends the whole saved transcript on every
request, so "keep the newest N tokens" recomputes from scratch each time and slides
forward a turn or two at a time. Measured on a 40-turn thread against an 8k window,
the eviction boundary moved on 12 of 40 turns, which means every few replies quietly
lost a little more of the conversation, llama-server's prefix cache was thrown away
each time the head of the prompt moved, and there was no such thing as a compaction
event to tell the user about.

Two changes make it discrete. The fit now reads back the boundary the thread last
compacted to, from the newest assistant turn's own persisted truncation, and reapplies
it before deciding anything, so nothing new is stored and it survives a restart. And
when the boundary does have to move, the trim takes a further ROLLING_COMPACTION_
HEADROOM_RATIO of the budget out (default 0.25) rather than skimming to the brim, so
the new boundary has room to stay put. Same 40-turn thread, same window: 4 compactions
over 60 turns instead of 14, keeping about 82 percent of the usable budget.

Both are gated on the prompt not already fitting, exactly as the recall reserve is. A
conversation inside its window is never evicted to satisfy either, and a stale boundary
from a branch that was rolled back cannot evict a chat that now fits.

The notice follows: it is shown when dropped_messages rises above the last turn that
reported it, so it appears once per compaction and stays quiet in between.

* Studio: pin that the compaction notice can never become conversation

The notice is a sidecar rendered from metadata, not a message, and the ways that
could quietly stop being true are all one careless edit away: moving it inside
MessagePrimitive.Parts would make it a content part, and everything that walks parts
would then replay it to the model, copy it and export it.

Asserts it renders as a sibling of the content parts, that neither the outbound
message builder nor the assistant replay serialiser reads the key it renders from
(bounded to those function bodies, since the streaming handler reads the same key
legitimately on the way in), that the markdown export does not mention it, and that
it is suppressed while editing so it cannot be typed into the textarea and saved
back as text.

* Studio: say which part is too long when nothing can make a turn fit

If the message just sent is itself bigger than the window, no amount of eviction
helps. The fit already handled that correctly, returning the conversation untouched
so the request reaches llama-server's normal context-length error, but what the user
was then told was actively misleading: the error reports the size of the WHOLE
conversation and advises shortening it, when the history has already been evicted and
the single message is the part that does not fit. Measured at a 4096-token window: a
5000-token message produces 'Message too long: 10290 tokens ... shorten the
conversation', and shortening it cannot possibly work.

The fit now returns a fits:false diagnosis instead of a bare None, carrying what the
conversation could not be reduced below and how much of that is the latest turn.
Every consumer already gated on fits, so this is inert wherever a truncation is
treated as a compaction; the streaming paths now forward it so the client can use it.

The toast reads the diagnosis and, when the latest turn alone exceeds the window,
says so with the numbers instead of offering advice that leads nowhere. The merge
drops the diagnosis once a later refit succeeds, by delete rather than by assigning
undefined, so an ordinary response keeps exactly the shape it had before.

* Studio: fix the review findings on the conversation archive

Fifteen items, each reproduced against the code before changing anything.

Correctness in the request path. The fit reported a successful recall-capable fit
whenever the result was under the prompt budget, but protected messages can stop the
trim reaching the reserve target, and the recall then went in anyway: reproduced at
ctx 8000, the fit accepted at 6900 and recall took the request to 8948, past the
window it had just been made to fit. Recall is now sized from the room the fit
actually obtained, and skipped when there is none. Separately, the sticky boundary
describes the original transcript, so re-applying it on a later tool-loop fit evicted
another boundary-sized block of live history: measured, a second fit dropped 28 of the
30 surviving messages instead of 14, and the summed count persisted an inflated
boundary for the next request. It is now spent after the first fit of a request.

Availability and privacy. enabled() trusted RAG_AVAILABLE, which only records that
import sqlite_vec worked; rag_available() exists because the native vec0 library it
loads is a separate file a venv can lack. On such a machine the fit held a recall
reserve back, evicting extra history, and then both the archive write and the recall
failed, so the user paid for content they never got. And a temporary chat is never
written to studio.db, yet the frontend still sends its thread_id and the request
carries no incognito flag, so its turns were archived to a scope no deletion flow
could reach. Archival now requires the thread to be persisted, which is the same rule
that keeps every archive reachable by a delete. Clear-history and project deletion
drop archives too; only DELETE /threads did.

Archive integrity. The document was committed before its chunks, so a failed chunk
write left an empty row marked completed that document_by_hash then skipped forever;
both now go in one transaction. The live-branch filter accepted a turn on its first
matching line, so editing only the assistant half kept serving the old answer; the
whole turn must be present. Retrieval fetched exactly k before that filter, so stale
turns could starve live ones and recall returned nothing; it over-fetches first. Tool
turns archived only the tool name, which cannot answer what was actually run, so a
bounded rendering of the arguments and the assistant text goes in as well.

The tool surface. Studio always sends an explicit enabled_tools array and has no
reason to name an internal tool, so the allowlist filter removed search_conversation
before the archive gate ran and the tool, plus the compaction nudge gated on it, never
appeared in a Studio chat. It now follows the archive rather than the allowlist. The
forced recall rendered a tool exchange even when the tool was absent from the
catalogue, which is the strict-template hazard the plain path avoids; it picks inline
in that case, and always inline for the final-answer request, which sends no tools at
all. A model-supplied top_k reached a slice as out[:-1] and returned nearly the whole
candidate pool, so it is clamped. Both retrieval tools now share the per-turn search
cap; only the knowledge-base one was counted.

The UI. A fits:false diagnosis is the fitter reporting it could NOT fit, so toasting
that older turns were removed was untrue and burned the once-per-thread flag a later
real compaction needed. The too-long advice compared the latest turn against the raw
context length rather than the prompt budget, so a 3,500-token message in a
4,096-token window was still told to start a new chat, which fails identically.

* Studio: fix the follow-up findings on the archive fixes

Four items, three of them about last round's own fixes.

Archived tool turns had become permanently unrecallable. render_turn now writes
'assistant called X: args' and 'tool result: ...' lines, while assistant-ui persists a
tool call as a structured tool-call content part that the transcript flattener dropped;
with every archived line required to appear in that transcript, no tool turn could ever
match. The transcript now flattens toolName, args and result, and the probe strips the
'assistant called <name>:' label, which is ours rather than the stored message's.

The branch probe compared only the first 160 normalized characters of each line, so an
edit to the tail of a long answer left the stale copy eligible. Ordinary lines are now
compared whole; tool results keep a prefix, since render_turn deliberately truncates
those and the archived copy is not meant to equal the stored one.

The forced recall sized itself by dividing the remaining budget by CHUNK_TOKENS, which
is an embedding-token limit rather than the chat template's cost, and prices none of
the wrappers around the injection. It is now recounted with the same tokenizer the fit
used and dropped if it overshoots, so the estimate can no longer eat the reply reserve.

An omitted top_k on search_conversation defaulted to the clamp ceiling of eight rather
than the configured recall default, so an ordinary search could return eight archived
turns into the protected current exchange that rolling truncation cannot evict.

* Studio: filter conversation recall to the active branch

A thread's stored rows are the whole message DAG. Retry and regenerate keep the
replaced response as a sibling on purpose, so filtering recall against the whole
thread cannot tell a live turn from one the user replaced, and an archived copy of
the abandoned response could be recalled into a branch where it never happened.

Filter against the messages the request was actually sent with instead, which is
one branch by construction, and hand the same branch to search_conversation so the
model cannot ask for what the forced recall refused. Falls back to the thread-wide
blob for a caller with no branch to offer.

Also retunes two respawn-refit fixtures whose windows no longer produced two
compactions after compaction started trimming a headroom margin below the budget.

* Studio: survive a delete mid-archive, and read the boundary off the active branch

Deleting a chat cancels its generation, but cancellation is cooperative and the
chunk-and-embed pass between the archive's liveness check and its commit does not
observe it. A delete landing in that window drops the thread's rows and sweeps its
scope before the commit puts rows back, leaving content the user deleted in a scope
no later delete can reach. Re-check after the commit and drop the scope: the delete
route removes rows first and sweeps archives last, so either order converges.

The sticky compaction boundary had the same thread-wide read as recall did. The
stored rows are the whole DAG ordered by creation time, so after a Retry the newest
assistant turn can be the sibling the user switched away from, and its boundary is
sized for history the active branch does not have. Resolve it against the request's
own messages instead.

* Studio: give the safetensors loop the same conversation-search guards

search_conversation is advertised by thread, not by backend: the tool selector is
shared, so a chat compacted under a GGUF model still offers it after the user
switches to a safetensors one. The safetensors loop had neither guard the GGUF loop
applies to it. It passed no active branch, so a search there fell back to the
thread-wide rows and could answer from a branch Retry left behind, and it capped
only search_knowledge_base, so paraphrased conversation searches could append
archived passages into the protected current exchange on every iteration until the
window failed. The shared set of capped retrieval tools now lives beside the cap.

The forced recall also derived its query from the loop conversation, which on a
later iteration can end with an internal user-role re-prompt rather than anything
the user wrote. It reads the request branch's own latest user turn instead.

* Studio: widen recall past an abandoned branch, and name whose turn overflowed

One over-fetch is not enough for the live-branch filter. Rewinding or retrying a
continuation that had already been compacted leaves enough stale turns to fill any
fixed candidate window, and the whole page is then rejected while the live match
sitting just below it is never examined, so recall reports nothing although the
answer is in the archive. Widen and re-ask instead, stopping as soon as there are
enough live hits, when the archive stops yielding candidates, or at a bound.

The irreducible-fit diagnosis also carried the size of the last message without
saying whose it was. A tool loop refits with the tool result appended, so that turn
is often output the user never wrote and cannot edit, and the client told them to
shorten it. It now reports the role, and the advice splits on it.

* Studio: shrink an over-budget recall, and keep truncated tool turns on their branch

The exact recount is the right gate, but dropping the whole recall when it fails is
the wrong response: with the shipped defaults a full top-K of long turns lands just
over the reserve once the wrappers are priced, which would disable the forced
retrieval on exactly the long conversations it exists for. Halve the number of turns
and re-ask instead, down to one, before giving up.

The live-branch probe keyed its prefix rule off the tool-result label, but a long
tool result is one appended string containing many newlines, so only its first line
carries that label. Continuation lines were compared in full, including the last one,
which is the only line the truncation marker is on and can therefore never appear in
a transcript: every archived tool turn over the cap was rejected as rolled back.
Key off the marker instead, which also covers truncated tool arguments and compares
the full text everywhere else.

* Studio: retire an edited turn's whole archived copy, not just the edited chunk

A turn longer than CHUNK_TOKENS is stored as several chunks of one document, and the
live-branch filter ran per chunk. Editing the second half of a long answer therefore
retired only the chunks carrying the edit, and an untouched earlier chunk of the same
retired turn stayed eligible on its own. The unit that was archived is the turn, which
is what the filter already claimed to enforce.

Validate every chunk of the candidate's document before admitting any of it, cached
per call since candidates from one turn share a document. A query failure falls back
to the per-chunk answer rather than failing the recall.

* Studio: do not hold a recall reserve back for a chat that is never archived

archive_turns refuses a thread with no saved messages, because a temporary chat must
not be persisted into a scope no deletion flow can reach. The reserve did not follow
that rule: it was granted whenever the RAG stack was available, so an incognito chat,
or an API client sending a thread_id without saving anything, paid a full 2,048-token
reserve for content that can never arrive. The fit subtracts the reserve from its trim
target, so that room is bought with evicted history. Measured on a 4K window: 15
tokens of conversation survived a compaction instead of 1,615.

Both now ask one predicate, can_archive, backed by an existence probe rather than a
row load.

* Studio: match archived turns in order, and delete them without sqlite-vec

Independent line membership accepts a turn whose lines were merely rearranged: every
probe still occurs somewhere, so the pre-edit ordering stayed eligible and would be
served back as what happened. Match the probes in order instead.

That only works if both sides agree on the order, and they did not: render_turn writes
a tool call before any assistant text on the same message and the result after it,
while both transcript builders wrote the text first, so a tool turn carrying both was
rejected outright. One flattener now lays out a message the way render_turn does, for
the request shape and the stored shape alike, and it offers both JSON spacings for a
stored call's arguments, whose object form is not the string the model emitted.

Deletion no longer depends on the optional native extension. An archive is only written
while vec0 loads, but the library can stop loading afterwards, and a delete that quietly
did nothing left a deleted conversation's turns on disk to answer again once it loaded.
The fallback removes the text-bearing rows over a metadata connection; the embedding
rows it cannot reach carry no text and resolve through tables that are gone.

* Studio: give the provider loops the branch, and keep the thread across a respawn

The provider tool loops take their catalogue from the same selector as the local ones,
so search_conversation is advertised there too once a thread has an archive, but the
loop passed no active branch and its searches fell back to the whole stored DAG, where
Retry keeps the response it replaced. This is the third loop to need the same wiring.

The respawn retry also dropped the thread. It re-enables context_overflow on purpose,
to refit an already-compacted prompt for a replacement window, so it is the one path
that deliberately compacts a second time: without the thread those extra evictions were
archived nowhere, nothing was recalled in their place, and the fit held back no reserve
and re-applied no boundary.

* Studio: keep the reply that follows a recall, and archive respawn refits

group_turns keeps an assistant tool call, its result and the reply that follows in one
group, and the archiver rejected any group containing one of our own injections. So on
a turn that forced a recall, the model's actual answer was thrown away with it: the
question was archived from its own group and the answer was not, and a later search
could find what was asked and never what was said. The injections come out now and the
rest of the turn stays; retrieved passages are still kept out of the index they came
from.

The two respawn refits inside the tool loop also evicted without archiving. They run
against a smaller replacement window, so those turns are simply gone otherwise. They
archive only, deliberately: no reserve is held back on that path, so injecting a recall
there is what would push the retry back over the window, and the next request can still
recall what this one archived.

* Studio: bound the branch check to one turn, and anchor only what recall injected

The branch check searched one flattened transcript, so a line an edit removed could be
supplied by any later message that happened to repeat the words. Short answers repeat
constantly: an archived "Should I deploy? / No" survived its answer being edited to
"Yes" because a later turn said "No", and the stale pair stayed recallable. The
transcript is now one normalised string per message, and a turn matches only if its
lines appear in order within a run of adjacent messages no longer than the turn itself.

Recall anchoring took the last two messages of the conversation. That is right for the
tool style, which appends a synthetic pair, and wrong for the inline style, which
appends nothing and rewrites the latest user message in place: it also pinned the
assistant turn before it, and with it an eviction unit the fit was entitled to drop,
which can fail a later iteration that would otherwise have fit. The injection now
reports exactly what it added or rewrote, and only that is anchored.

* Studio: validate a turn's chunks against one run of the branch, not each on its own

The chunks of an archived turn are consecutive slices of a single rendering, and each was
checked for itself. That let a turn be reassembled out of parts that never sat together:
the head matching the question and the answer it has now, the tail matching some later
message that happens to repeat the passage an edit removed, so the stale association
stayed recallable. Reproduced before the change, with each chunk passing on its own.

All chunks of the document must now be found within one run of adjacent messages, bounded
by the document's own line count, which is at least the number of messages the turn was
rendered from, so a turn that really is still there always fits.

* Studio: budget conversation searches and disambiguate identical replies

search_conversation clamped the model's top_k only against a fixed ceiling of 8.
Eight chunks is roughly 4,000 tokens once wrapped, and the result lands in the
current tool exchange, which rolling truncation protects and cannot evict, so on
a small context the search itself made the turn unsendable. The GGUF loop now
passes the room the window actually has left, the tool clamps top_k by it, and a
search with no room says so instead of returning a result that cannot be sent.

The sticky compaction boundary took the newest on-branch assistant row, but the
branch check is textual, so two siblings whose replies read the same ("Done.")
are indistinguishable from there, and Retry is exactly what produces them.
Taking the first match applied a boundary measured on a different, deeper branch
and evicted live history. Where the text cannot separate them, the smallest
boundary is now used: too small costs one more compaction, too large loses turns
the branch still has.

* Studio: price a conversation search against the whole prompt

Three gaps in the budget the previous commit introduced.

The GGUF loop subtracted an estimate of the messages alone, leaving the tool
catalogue out of the prompt entirely. A large catalogue is thousands of tokens,
so the request could already be near its budget while the search was still told
there was room for several 500-token chunks. The preflight already prices the
request exactly, catalogue and template included, so the difference between that
count and the estimate of the same messages is carried forward and the estimate
only covers what the loop appends after it.

The safetensors loop advertises and executes search_conversation for a thread
compacted under a GGUF model, but forwarded no budget, so the clamp in the tool
was skipped there. It now passes the room this model has left. With no known
context length the argument is omitted rather than sent as zero, which would
refuse every search.

The sticky boundary filtered candidates with a substring test, which is right
for archived chunks and wrong for whole messages: an abandoned "Done" rode in on
a live "Not done yet", and having no live twin it then decided the boundary
alone. Where any candidate matches a live message exactly, only the exact ones
are considered; where none does, the old behaviour stands, since a stored row is
not always byte-identical to what the client re-sends.

* Studio: keep the overflow diagnosis on the tool path

The irreducible-overflow branch counted the newest message on its own, and a
tool loop reaches that branch with a tool result last. A tool result by itself
is not a conversation: templates that require it to follow its assistant tool
call refuse to render one, and the exception escaped the fit entirely, so the
caller fell back to the untrimmed request and the client was told nothing at
all. That is the one path this diagnosis was added for.

The count is now attempted and falls back to the estimator when the template
refuses. An approximate number is worth more here than a diagnosis that never
arrives.

* Studio: budget the conversation search in the provider loop

The third tool loop forwarded the active branch but no budget, so the clamp in
the tool was skipped there and a model-chosen top_k of 8 could append roughly 4K
tokens to a prompt this loop replays on its next call.

Studio knows no window for an external model: the request carries no context
length and there is no registry to look one up in, and a custom
OpenAI-compatible endpoint can be a small local server. So rather than a
measured budget, this path spends no more than one ordinary recall's worth,
which is the same amount the compaction turn itself is sized for.

* Studio: bound an archived turn to its own messages, and archive it once

Two problems in the archive.

The whole-document branch check bounded the run by the number of LINES the turn
produced, and gave each chunk its own start. A turn is two or three messages
however long it is, so a hundred-line answer got a hundred-message window, and
the tail of an edited answer could be satisfied by a message well outside the
turn. The run is now bounded by the messages the turn was rendered from, counted
from the labels render_turn writes, and the chunks are scanned as one pass: each
continues where the previous one stopped rather than restarting. The cursor
inside a message is deliberately not carried over, because chunks overlap and a
continuation chunk repeats the tail of the one before it.

The hash check that skips an already-archived turn ran long before the insert,
with the embedding pass in between, and the index on (scope, sha256) is not
unique. Two generations compacting the same thread both cleared it and both
wrote, so the turn was stored twice and its copies took two of the few recall
slots. The check is now repeated under a write lock immediately before the
insert. Reproduced with two concurrent archive passes: two documents before, one
after.

* Studio: record how many messages an archived turn came from

The run an archived turn is allowed to occupy on the branch was bounded by
counting the role labels in its rendered text. That counts lines the user wrote
as well as the ones the renderer did: a pasted chat log carries lines that look
exactly the same, and each one widens the run by a message, which is enough for
the message after an edited turn to supply the passage the edit removed.

The group's size is now recorded on the document when it is archived, in a
nullable column added the way the other lazy upgrades are. Archives written
before this have NULL and fall back to the label count, so nothing needs
backfilling to keep working.

* Studio: persist a compaction boundary the next request can use

The boundary was read back from dropped_messages, which counts what each fit
removed from the conversation in front of it. The tool loop refits on every
iteration and the client sums those counts, so a long agent run added the tool
exchanges the turn itself created, and the next request applied the total to its
saved transcript. Reproduced with six tool calls on a 4K window: three fits of 4
summed to 12, on a branch that only ever had 4 evictable messages.

The boundary is now carried separately, measured against the messages the
request was sent with, so it is absolute and re-sending it cannot advance it.
The client keeps the latest value rather than summing, and turns saved before
this fall back to the dropped count, which is the same number for a turn that
fit once.

The recall reserve is also dropped once archiving has failed. sqlite-vec can be
present and the thread saved while the embedder cannot start: archive_turns
swallows that and recall injects nothing, so the room the fit held back was pure
loss on every compaction, and the failure mode forgot more history than having
the feature off. A failed write marks the archive degraded and the next
successful one clears it.

* Studio: count the boundary past the system prompt, and notice it

Two faults in the boundary added in the previous commit.

It stopped at the first message still present, and a Studio request always
starts with a system prompt that a fit never evicts, so every compaction
recorded a boundary of zero. That is the same as having no boundary: the next
request would move the eviction point again and invalidate the prefix cache on
every turn. Instruction messages are now skipped rather than treated as the
front of the branch, and the newest turn is excluded because it is never evicted
and an inline recall rewrites it in place.

The compaction notice still keyed off dropped_messages, which is the
accumulated count. A tool-heavy turn reporting 12 while the boundary moved to 4
set a high-water mark that silenced the next two real advances. It now reads the
boundary, falling back to the dropped count only for turns saved before the
boundary was recorded.

* Studio: tighten comments in the conversation archive and rolling context window

* Studio: give the pasted-text import an extension on this branch too

The node test runner cannot resolve an extensionless relative import, so
delete-chat-files-preference fails to load the preferences store. main fixed
this in d43892ea7; this branch predates it, and the same line is on the branch
below it. The change is identical to main's, so it disappears on merge.

* Studio: size a conversation search by what it renders to

CHUNK_TOKENS is what the chunker aims at, not what a chunk weighs: chunks
overlap, the chunker's tokenizer is not the model's, and the rendered block adds
markup, source metadata and the tool framing around it. Dividing the budget by
CHUNK_TOKENS therefore permitted a result far larger than the room measured, and
it lands in the current exchange, which the rolling window cannot evict.
Measured on a 500-token budget: one chunk came back at 1,256 estimated tokens.
The count is now halved until the rendered result fits, the same backoff the
forced recall uses, and a single chunk that still does not fit is refused.

The late archive cleanup also spared a recreated thread. DELETE removes the rows,
awaits the sandbox pass, and only then sweeps the archive; another tab can POST
the same id in that window and its generation can archive turns under it. The
sandbox pass re-checks for exactly that, and this now does too, so the recreated
chat keeps its memory.

* Studio: keep the configured default when top_k is omitted

Budgeting an omitted top_k by dividing the whole budget treated the room as a
target rather than a cap: on a 128K chat with most of its window free it asked
the archive for 200 passages, past the configured default of 4 and past the
ceiling of 8 that the model's own value is held to. The default now applies as
before, with the room and the ceiling capping it.

* Studio: cut the boundary at the newest user turn

Excluding only the branch's last message assumed the newest user turn is last,
and a continued assistant message puts a prefill after it. The user turn then
stayed in the identity scan, and since inline recall rewrites it into a new dict
it read as evicted, inflating the boundary by one and costing the next request a
live message. The scan now stops at the newest user turn, which is what is
protected in any case.

* Studio: require an archived turn to end where the live message does

Editing a reply by keeping it and adding to it ("No" becoming "No, correction:
yes") left every probe matching, so the pre-edit copy stayed eligible and a
search could return "No" as the answer with the correction nowhere in it. The
scan now reports where it finished inside the message, and the turn is accepted
only if it reaches the end of it.

Tool results and tool arguments are exempt: render_turn cuts them, so their
probes are prefixes by design and demanding the end would retire every turn that
carried one.

* Studio: check both ends of an archived turn, and charge dense text properly

An edit that keeps the old text and adds to it leaves every probe matching
whichever side it adds on. The end anchor caught "No" becoming "No, correction:
yes"; it did not catch "Correction: no", which ends exactly where the archived
copy does. The scan now also reports where the run into the final message began,
and the turn is accepted only if it covers that message end to end. render_turn
only ever cuts the tail, so the start is required unconditionally.

A conversation search sized itself with the shared estimator, which charges four
characters per token. That is an English rule: CJK and emoji run closer to one
token per character, so a result could be accepted at a quarter of its real size
and then land in the current tool exchange, which the window cannot evict. No
exact counter is reachable there, the provider loop having no tokenizer at all,
so non-ASCII characters are charged one token each and the rest at the usual
rate. Measured on 320 CJK characters: 87 tokens claimed, 320 charged.

* Studio: earn the compaction headroom, and keep this turn's tool results

The headroom cut every overflowing prompt to about 75% of its budget, including
on paths that can never put the boundary back: an incognito chat, an API request
with no persisted thread, or one whose turns are not saved. There the deeper cut
buys nothing and cannot be recalled, so it is simply less history than plain
eviction would have kept, and turning the archive off did not restore the old
behaviour. The fit now takes the headroom only when the caller can restore the
boundary next request.

The branch the archive is filtered against was the messages the client sent. A
long agent run evicts, and archives, tool exchanges it created itself, and those
were then refused as an abandoned branch: the model could not search back a tool
result it still needed to answer. The branch is now accumulated across the
request, so this turn's own exchanges count as live, while a sibling response
from a Retry still does not.

* Studio: read the branch as replies when restoring a boundary

The rows the rolling window checks are assistant replies, but the branch they
were checked against was every message of it, flattened without roles. An
abandoned "Done" left by a Retry therefore rode in on a live user message that
merely contains it ("not done yet I think"), and with no exact match to prefer,
its much larger boundary was applied to a branch that never had that reply.

The branch representation now takes a role filter and the boundary asks for
assistant messages only. A branch with no reply of its own has no boundary to
restore, so it reports none rather than falling back to the unfiltered check.

* Studio: account for every message an archived turn claims

The anchors covered the final message only, so an edit to the question
underneath it left the archived copy eligible, whichever side the text was added
on. The scan now checks each message the run touches: it must be matched from
the first character, and nothing may be left over when the run moves on.

Tool calls are exempt, and the exemption covers the whole message. The store
keeps a call as a structured part, so the live text carries the tool name and
both spellings of the arguments, spaced and compact, while the archived copy has
one line of one of them. Nothing there lines up character for character, and
demanding it would retire every tool turn in the archive.

* Studio: re-embed archived turns after a model change, and batch the first pass

Deduplication was by hash alone, but dense search only reads documents whose
recorded embedder matches the query's. A turn archived under the previous model
was therefore skipped on every later compaction while being invisible to every
paraphrased search, permanently. The check now compares the identity too and
replaces the copy, which is what ingestion does with a re-uploaded file.

Embedding also ran once per evicted turn. A first compaction of a long chat sent
dozens of one-item jobs back to back, and both backends serialise them, so the
reply waited for all of them. The turns are chunked first and embedded in one
pass, with the per-turn write and its duplicate check unchanged. Measured on 40
turns: 40 embedding calls before, 1 after.

* Studio: re-pin the permission read the deep-research contract lost to #8686

The contract still asserted the literal `const permissionMode = loadPermissionMode();`,
which #8686 removed when it put a chat-scoped override in front of that read. This branch
already carries the refactored store, so Repo tests (CPU) cannot pass here or on anything
stacked above it.

Ported main's replacement rather than dropping the check. The read is still the contract,
and it is scoped to the deep-research setter so it cannot be satisfied by the initial-state
constant, which is a different property and would keep passing if the read were deleted.

* Studio: re-derive the compaction boundary by position, not by count

`boundary_messages` is an absolute count against the transcript it was counted
on. Delete an already-evicted prompt from a compacted thread and the front of
that transcript gets shorter, but the count is replayed unchanged, so the cut
lands that many messages too deep and evicts turns that are still live.

Reproduced on `fit_rolling_context` with a 30-turn chat at ctx 2000: a fresh fit
keeps from turn 16; deleting one already-evicted pair and replaying the stored
count keeps from turn 17, one live turn lost. The loss is then baked in, since
the next turn records the boundary against the shortened transcript.

Stamp `boundary_anchor`, the text of the first message the fit KEPT, next to the
count at all five sites that write it, and look it up on the next request's own
branch. The anchor is only ever allowed to make the boundary SHALLOWER, so a
stale, ambiguous or repeated anchor costs one extra compaction and can never
evict a live turn. Rows written before the key exists behave exactly as today.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: say what the archive kill switch does and does not cover

RAG_CONVERSATION_ARCHIVE=0 stops indexing, recall and the recall reserve. It
does not turn the rolling window back into what it was before it learned to
compact in discrete events: the headroom and the sticky boundary belong to the
window, and the headroom has its own knob. "Off restores plain eviction" read
as though it covered all of it.

Comment only. Gating the window's own behaviour on this flag would make a host
with no sqlite-vec silently compact differently from one that has it.

* Studio: charge dense text honestly, and reject an edit inside an archived turn

Two fixes.

The safetensors search budget priced what is already in the prompt at four
characters per token. That is about right for English and roughly half the truth
for CJK and emoji: measured on an 81-message Japanese chat, 1295 estimated
against 2737 real, reporting 1777 tokens of room where 335 remained. The tool
result is already charged a token per non-ASCII character; this is the same rule
applied to the spend, which was the missing half. That path runs no rolling fit,
so nothing downstream recovers once the exchange lands. The GGUF site takes it
only on its fallback leg, since the fit's exact tokenizer count needs no
correction, and eviction keeps the flat estimate: making eviction pessimistic
would drop history a request could have kept.

The probe scan anchored the start and the end of a run but not the middle, so a
correction inserted between two archived lines matched both probes with the new
line sitting unexamined in the gap, and the pre-edit turn stayed recallable.
Verified on "A\nB" becoming "A\ncorrection\nB". A gap is now allowed only where
it is a label render_turn wrote and the probe had therefore stripped, since a
pasted chat log legitimately carries its own "user:" lines, and the tool-call
exemption still covers the rest of its message.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: end the tool-call anchor exemption where the call ends

A stored tool call cannot line up character for character with the live text,
because the store keeps arguments as an object and offers both JSON spellings,
so the cursor after one is not exact and the character anchors have to relax.
That relaxation was then held for the rest of the message.

An assistant turn carrying both a call and text therefore stayed matched after a
correction was appended to the text: "old answer" becoming "old answer,
correction: new answer" left the archived pre-edit turn eligible for recall.

Cleared once an ordinary text probe matches, since the cursor is exact again from
there. Verified on that exact pair: intact still eligible, edited no longer.

* Studio: count a video part as media in the rolling preflight

`_inject_video_part` writes llama.cpp's own `input_video` part into the same
message list the fit is then handed, and the media predicate did not list it.

Video prompts therefore ran the rolling preflight, which is skipped for media
precisely because `/apply-template` token counting does not include the sampled
video tokens. The prompt could be certified as fitting when it does not, or lose
history it never needed to lose, and still fail with context_length_exceeded.

* Studio: keep which tool was called in the branch probe

`render_turn` writes "assistant called <name>: <args>", and the whole label was
stripped before probing, name included, so only the arguments decided. A retry
that kept the arguments and the result and changed only the tool left the
archived pre-edit turn matching: verified on the same arguments moving from
terminal to python, which stayed eligible and could be recalled as though the
old call had happened on this branch.

The name is now its own probe, ahead of the arguments, which is the order
`_probe_text` renders a live call in. The label itself still has to go, since it
exists only in the archived copy, but the name does not: the live text carries
it. `render_turn`'s "tool" fallback for a nameless call is excluded, because the
live text has no name there either.

Not an exact anchor: the name is matched as a substring of the message, so a
tool whose name appears inside the arguments of another can still match. This
closes the ordinary retry case rather than every construction of it.

* Studio: do not archive a search the model asked for

`_is_injected` recognises the ids this feature and the RAG auto-inject generate,
and a search the MODEL issues carries neither: the parser gives it an ordinary
`call_N` id. Both the call and the passages it retrieved were therefore indexed
as fresh conversation, and a second search archived the first one's output inside
its own, one nesting level per distinct search, each copy competing for the four
recall slots.

Dropped by tool NAME instead, reusing `RAG_SEARCH_TOOLS`, which both loops
already share. Only the retrieval parts go: a reply that follows a search is real
conversation and is still archived, and an assistant message carrying a retrieval
call beside an ordinary one keeps the ordinary one and its result.

* Delete the archived conversation even when its thread id comes back

The late archive sweep skipped the whole scope when another tab had recreated
the id, which spared the recreated chat's memory but kept the deleted
conversation's too. The scope is keyed by thread id alone, so those turns stayed
recallable in the new chat with nothing left to sweep them, and the endpoint
reported success.

Take a cutoff before the rows go and bound the delete by created_at when the id
has come back: everything archived before the delete was accepted belongs to the
deleted conversation, everything after to the new one. Applied at DELETE
/threads, project delete and clear all, and on the no-vec0 path so it cannot
delete more than the vec0 path does.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: tighten rolling context window comments

* Stop folding case when validating an archived turn against its branch

A turn corrected only in capitalisation still matched its archived copy, so the
branch filter kept the pre-edit document live and a recall could answer with the
spelling the user had just fixed. Nothing else retires it: the edit changes the
digest, so the corrected turn is written as a new document and the branch filter
is the only thing that could have dropped the old one.

Trim both ends rather than only the right. Keeping leading whitespace looks
tighter and is worse, since render_turn strips the whole message and a live turn
opening on a space or a newline then starts its run at a non-zero offset and is
retired outright. Measured: '   hello there' and a pasted block opening on a
newline both went from live to retired. Indentation-only edits stay tolerated in
both directions as a result, because probes are matched per line and as
substrings; closing that needs the splitter to carry offsets.

Also fixes an unrelated retirement found while measuring this. When
render_turn's 4000 character cut lands exactly on a newline the marker becomes a
line of its own, strips to empty, and was dropped along with the truncation flag,
so the last real probe read as complete and an unedited over-cap tool result was
retired. Measured on a 900 line result: no query could return it.

Note for anyone testing an upgrade in place: boundary anchors written by an
earlier build of this branch were stored folded, so they will not match until the
next compaction rewrites them. The anchor can only ever move the boundary
shallower, so the cost is one extra compaction.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Let search_conversation through the Anthropic gate, and probe a tool call before its answer

Adding the schema to ALL_TOOLS made the Anthropic selector pick it while
_ANTHROPIC_UNPROMPTED_SAFE_TOOLS still listed only web_search and
search_knowledge_base, so the pre-switch guard classified a read-only tool as
confirmation-gated. Measured: enable_tools with enabled_tools
['search_conversation'] returned 400 on auto and on the omitted default, with
the terminal/python message, even though is_potentially_unsafe_tool_call marks
it always safe and that set is documented as mirroring it. Pre-PR the same
request was served, so this is a regression the PR introduced.

Separately, _probe_text bucketed a whole message as call, text, result, which
only holds while the text came before the call. A persisted row whose tool call
is followed by the model's final answer, the ordinary agent turn, renders as
call, result, answer, so _scan_probes advanced past the answer to find the
result and could not find it again. Measured end to end on the persisted-row
transcript: recall returned the user's question alone and the document holding
the answer was filtered out. The buckets now flush when text arrives after a
call, the same rule the replay serializer uses; text written before a call still
rides ahead of it.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Record the transcript span, and price the recall budget exactly when nothing was trimmed

archive_messages bounds the branch check's run, and it was recording the number
of messages ARCHIVED rather than the span of the turn. The two differ whenever
_archivable drops something: an assistant batch that called search_conversation
alongside an ordinary tool archives three messages while the live transcript
holds four. Measured, validation failed at three and passed at four, so a
perfectly valid ordinary-tool exchange and the answer that followed were
rejected as off-branch and could never be recalled.

fit_rolling_context returns None when it drops nothing, so a prompt that simply
fits, after a context-length increase or on a shorter branch, left the recall
budget to a character estimate that cannot see the template's own framing.
Measured on a request whose real prompt is about 2800 tokens of a 3584 budget:
the budget came back 3512, nearly the whole window, so the recall could append a
passage that does not fit and the next iteration cannot evict it again, since the
current tool exchange is protected. Priced exactly instead, once per request and
only when the model actually reaches for a retrieval tool on a request that did
not truncate. A failure falls back to the old estimate.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Price the recall result with the model tokenizer, recount every search, and anchor on tool calls

Count the search result with the caller's own token counter instead of the ASCII
heuristic, recompute the exact prompt count on every search rather than once per
turn, and take the boundary anchor from a message's tool calls when it has no text,
so a tool-call message no longer records an empty anchor and disables the rebase.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Widen a reused turn's window to its longest span, and price every search exactly

A three-message tool exchange and a four-message batch holding the same exchange render
identically, so the second is skipped as a duplicate and inherited the shorter window,
which the branch check then used to reject it. Grow the stored span instead, upwards
only. Count the prompt exactly on every conversation search rather than only when the
fit dropped nothing.

* Let an overlapping chunk restart far enough back to cover its carried messages

A chunk's leading overlap is the previous chunk's tail, and it can carry a whole short
message with it, so resuming the scan at the message the previous chunk finished in
could never match it and retired unedited turns as off-branch. The forward position is
still tried first and the walk back stops at the previous chunk's own opening.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Widen the span on the locked duplicate path as well

Both turns can arrive in one compaction: the pre-check clears both before either is
written, the shorter is written first, and the longer then met the re-check under the
write lock and left the window at the shorter figure.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Bound the persisted anchor, and spend any room at all on one recall attempt

The anchor rides in every truncation event and every assistant turn's metadata while the
boundary stays sticky, so a large pasted message was copied across the thread; cap it at
a head on both sides, where the read side already clamps only shallower. And a budget
below one chunk is not no budget: CHUNK_TOKENS is a ceiling, not the size of a turn, so
try one and let the exact recount reject it.

---------

Co-authored-by: alkinun <alkinunl@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothshared@gmail.com>
2026-08-19 17:40:23 -07:00

3975 lines
142 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""SQLite storage for training run history and metrics.
Like auth/storage.py (module-level functions, raw sqlite3, per-function
connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes.
"""
import hashlib
import json
import logging
import os
import platform
import re
import shutil
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
logger = logging.getLogger(__name__)
from typing import Any, Iterable, Optional
from utils.paths import (
ensure_dir,
project_workspaces_root,
studio_db_path,
)
from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root
from utils.paths.scan_folder_health import is_readable_dir
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
from utils.training_runs import extract_project_name
def _extract_project_name_from_config_json(config_json: Optional[str]) -> Optional[str]:
if not config_json:
return None
try:
return extract_project_name(json.loads(config_json))
except (json.JSONDecodeError, TypeError):
return None
def _denied_path_prefixes() -> list[str]:
"""Platform-aware denylist of system directories."""
system = platform.system()
if system == "Linux":
return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"]
if system == "Darwin":
# macOS realpath() resolves /etc -> /private/etc etc; include /private variants.
return [
"/System",
"/Library",
"/dev",
"/etc",
"/private/etc",
"/tmp",
"/private/tmp",
"/var",
"/private/var",
]
if system == "Windows":
win = os.environ.get("SystemRoot", r"C:\Windows")
pf = os.environ.get("ProgramFiles", r"C:\Program Files")
pf86 = os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)")
return [os.path.normcase(p) for p in [win, pf, pf86]]
return []
def is_denied_system_path(path: str) -> bool:
"""True if *path* is, or descends from, a denied system directory.
Mirrors the denylist add_scan_folder() enforces at registration so the
browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds
a broad root (a Windows drive root C:\\ or a legacy-registered / root). The
/run carve-out keeps Linux removable-media mounts browseable. Expects an
already-resolved (realpath) path so symlinks cannot escape into a denied subtree.
"""
is_win = platform.system() == "Windows"
check = os.path.normcase(path) if is_win else path
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
return True
return False
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
def contains_sensitive_path_component(path: str) -> bool:
return _contains_sensitive_path_component(path)
_schema_lock = threading.Lock()
_schema_ready = False
_SQLITE_IN_CHUNK_SIZE = 900
_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",)
_CHAT_ATTACHMENT_INVENTORY_VERSION = 1
def _project_slug(name: str) -> str:
slug = re.sub(r"[^A-Za-z0-9._-]+", "-", name.strip()).strip(".-_")
return slug[:48] or "project"
def _default_project_root(project: dict) -> str:
project_id = str(project["id"])
suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project"
folder_name = f"{_project_slug(str(project.get('name') or 'Project'))}-{suffix}"
return str(project_workspaces_root() / folder_name)
class ProjectWorkspaceError(OSError):
"""Raised when a project's workspace folder cannot be created.
Tagged, and carrying the folder, so a caller can name it. The same upsert
also touches the database directory, and that is a different path with a
different fix.
"""
def __init__(self, path: str, cause: OSError):
super().__init__(str(cause))
self.path = path
def _ensure_project_workspace(root_path: str) -> str:
root = Path(root_path).expanduser()
try:
root_resolved = ensure_dir(root).resolve()
for subdir in _PROJECT_WORKSPACE_SUBDIRS:
ensure_dir(root_resolved / subdir)
except OSError as exc:
raise ProjectWorkspaceError(str(root), exc) from exc
return str(root_resolved)
def sandbox_is_referenced_elsewhere(
session_id: str, exclude_thread_id: "str | None" = None
) -> bool:
"""Whether a surviving chat still shows file cards for this sandbox.
Forking clones the message content verbatim, so the fork's cards keep the
source chat's session id. Deleting the source's files would leave those
cards downloading nothing, in a chat the user did not delete.
"""
if not session_id:
return False
conn = get_connection()
try:
# The LIKE only narrows, on the id as JSON writes it, so a quote or a
# backslash is still found. Every hit is parsed and the id has to be a
# sessionId value: a short one is a substring of ordinary prose.
escaped = json.dumps(session_id)[1:-1]
rows = conn.execute(
"""
SELECT content_json FROM chat_messages
WHERE (? IS NULL OR thread_id != ?) AND content_json LIKE ? ESCAPE '\\'
""",
(exclude_thread_id, exclude_thread_id, f"%{_like_escape(escaped)}%"),
)
for row in rows:
if _mentions_session(row["content_json"], session_id):
return True
return False
except sqlite3.Error:
# A locked database is not an answer, and every caller reads False as
# "nothing shows these files any more" before deleting them. Kept, so
# the worst case is a folder collected on the next delete.
logger.warning("Could not check references for sandbox %s; keeping it", session_id)
return True
finally:
conn.close()
def _mentions_session(content_json: str, session_id: str) -> bool:
"""Whether this message's content names *session_id* as a sandbox."""
try:
content = json.loads(content_json)
except (TypeError, ValueError):
return False
stack = [content]
while stack:
node = stack.pop()
if isinstance(node, dict):
if node.get("sessionId") == session_id:
return True
stack.extend(node.values())
elif isinstance(node, list):
stack.extend(node)
return False
def _like_escape(value: str) -> str:
for char in ("\\", "%", "_"):
value = value.replace(char, "\\" + char)
return value
def delete_project_workspace(project: dict) -> None:
"""Remove a deleted project's workspace directory.
Separate from the row delete so the caller can stop the tool calls running
in there first: pulling the working directory out from under a live
subprocess is how a half-written file ends up outside any project.
"""
_delete_project_workspace(project)
def _delete_project_workspace(project: dict) -> None:
root_path = project.get("rootPath")
if not root_path:
return
root = Path(root_path).expanduser()
try:
root_resolved = root.resolve(strict = False)
except (OSError, RuntimeError, ValueError):
logger.warning("Skipping project workspace delete for invalid path %r", root_path)
return
project_id = str(project["id"])
suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project"
if not root_resolved.name.endswith(f"-{suffix}"):
logger.warning(
"Skipping project workspace delete for unexpected project path %s",
root_resolved,
)
return
if root_resolved.parent == root_resolved or root_resolved == Path.home().resolve():
logger.warning(
"Skipping project workspace delete for unsafe project path %s",
root_resolved,
)
return
check = (
os.path.normcase(str(root_resolved))
if platform.system() == "Windows"
else str(root_resolved)
)
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
logger.warning(
"Skipping project workspace delete under denied path %s",
root_resolved,
)
return
if not root_resolved.exists():
return
if root_resolved.is_symlink() or not root_resolved.is_dir():
logger.warning(
"Skipping project workspace delete for non-directory path %s",
root_resolved,
)
return
shutil.rmtree(root_resolved)
def delete_chat_project_workspace(project: dict) -> None:
"""Backward-compatible name for callers predating sandbox-aware cleanup."""
delete_project_workspace(project)
def _ensure_schema(conn: sqlite3.Connection) -> None:
"""Create tables and indexes if they don't exist. Called once per process."""
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_runs (
id TEXT NOT NULL PRIMARY KEY,
status TEXT NOT NULL DEFAULT 'running',
model_name TEXT NOT NULL,
dataset_name TEXT NOT NULL,
config_json TEXT NOT NULL,
started_at TEXT NOT NULL,
ended_at TEXT,
total_steps INTEGER,
final_step INTEGER,
final_loss REAL,
output_dir TEXT,
error_message TEXT,
duration_seconds REAL,
loss_sparkline TEXT,
display_name TEXT,
resume_blocked INTEGER NOT NULL DEFAULT 0,
resumed_from_run_id TEXT
)
"""
)
existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()}
if "display_name" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
if "resume_blocked" not in existing_cols:
conn.execute(
"ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0"
)
# Nullable, so older rows stay NULL and fall back to the output_dir heuristic.
if "resumed_from_run_id" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN resumed_from_run_id TEXT")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_metrics (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL REFERENCES training_runs(id) ON DELETE CASCADE,
step INTEGER NOT NULL,
loss REAL,
learning_rate REAL,
grad_norm REAL,
eval_loss REAL,
epoch REAL,
num_tokens INTEGER,
elapsed_seconds REAL,
UNIQUE(run_id, step)
)
"""
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)")
# Windows: COLLATE NOCASE so C:\Models and c:\models dedup; elsewhere BINARY keeps them distinct.
collation = "COLLATE NOCASE" if platform.system() == "Windows" else ""
conn.execute(
f"""
CREATE TABLE IF NOT EXISTS scan_folders (
id INTEGER PRIMARY KEY AUTOINCREMENT,
path TEXT NOT NULL UNIQUE {collation},
created_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_projects (
id TEXT NOT NULL PRIMARY KEY,
name TEXT NOT NULL,
instructions TEXT,
root_path TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)
"""
)
chat_project_cols = {
row[1] for row in conn.execute("PRAGMA table_info(chat_projects)").fetchall()
}
if "root_path" not in chat_project_cols:
conn.execute("ALTER TABLE chat_projects ADD COLUMN root_path TEXT")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_projects_archived_updated_at ON chat_projects(archived, updated_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_threads (
id TEXT NOT NULL PRIMARY KEY,
title TEXT NOT NULL,
model_type TEXT NOT NULL,
model_id TEXT,
pair_id TEXT,
project_id TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER,
openai_code_exec_container_id TEXT,
anthropic_code_exec_container_id TEXT,
forked_from_thread_id TEXT,
forked_from_message_id TEXT,
settings_json TEXT,
FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE
)
"""
)
chat_thread_cols = {
row[1] for row in conn.execute("PRAGMA table_info(chat_threads)").fetchall()
}
if "settings_json" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN settings_json TEXT")
# Orders one writer's snapshot writes against its own earlier ones. A tab closing
# sends its last edit keepalive, which can overtake a PATCH already accepted by the
# server, and no client-side cancel reaches a handler that is already running: the
# older write has to be refused here. Scoped to the writer that sent it, so two
# browsers are never ordered against each other; see write_chat_thread_settings.
# {writer id: highest seq seen from it}. One writer and one seq is not enough: a
# write from another tab overwrites them, and the delayed request the ordering exists
# to refuse is then compared against a watermark that is no longer its own.
if "settings_seqs" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN settings_seqs TEXT")
if "project_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN project_id TEXT")
if "openai_code_exec_container_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT")
if "anthropic_code_exec_container_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT")
if "forked_from_thread_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT")
if "forked_from_message_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT")
if "updated_at" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER")
# Floor at created_at: forked threads copy older ancestor messages, so the fork's creation time wins.
conn.execute(
"""
UPDATE chat_threads SET updated_at = MAX(
COALESCE(
(
SELECT MAX(m.created_at) FROM chat_messages m
WHERE m.thread_id = chat_threads.id
),
created_at
),
created_at
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_thread_tombstones (
id TEXT NOT NULL PRIMARY KEY,
deleted_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_clear_operations (
id TEXT NOT NULL PRIMARY KEY,
active_research_run_ids_json TEXT NOT NULL,
deleted_thread_ids_json TEXT NOT NULL DEFAULT '[]',
cleared_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
chat_clear_operation_cols = {
row[1] for row in conn.execute("PRAGMA table_info(chat_clear_operations)").fetchall()
}
if "deleted_thread_ids_json" not in chat_clear_operation_cols:
conn.execute(
"ALTER TABLE chat_clear_operations ADD COLUMN deleted_thread_ids_json TEXT NOT NULL DEFAULT '[]'"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
id TEXT NOT NULL PRIMARY KEY,
thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
parent_id TEXT,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
attachments_json TEXT,
metadata_json TEXT,
created_at INTEGER NOT NULL
)
"""
)
tombstone_schema = """
CREATE TABLE chat_attachment_tombstones (
thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
message_id TEXT NOT NULL,
attachment_id TEXT NOT NULL,
deleted_at INTEGER NOT NULL,
PRIMARY KEY(thread_id, message_id, attachment_id)
) WITHOUT ROWID
"""
tombstone_table = conn.execute(
"""
SELECT 1 FROM sqlite_master
WHERE type = 'table' AND name = 'chat_attachment_tombstones'
"""
).fetchone()
if tombstone_table is None:
conn.execute(tombstone_schema)
else:
tombstone_columns = {
row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)")
}
tombstone_fk_targets = {
row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)")
}
if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets:
# The first implementation cascaded through chat_messages, which erased deletion knowledge
# during pruneMissing. Rebuild once, retaining every tombstone whose thread still exists.
conn.execute("SAVEPOINT migrate_chat_attachment_tombstones")
try:
conn.execute(
"ALTER TABLE chat_attachment_tombstones "
"RENAME TO chat_attachment_tombstones_legacy"
)
conn.execute(tombstone_schema)
if "thread_id" in tombstone_columns:
conn.execute(
"""
INSERT OR IGNORE INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
SELECT legacy.thread_id, legacy.message_id,
legacy.attachment_id, legacy.deleted_at
FROM chat_attachment_tombstones_legacy legacy
JOIN chat_threads thread ON thread.id = legacy.thread_id
"""
)
else:
conn.execute(
"""
INSERT OR IGNORE INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
SELECT message.thread_id, legacy.message_id,
legacy.attachment_id, legacy.deleted_at
FROM chat_attachment_tombstones_legacy legacy
JOIN chat_messages message ON message.id = legacy.message_id
"""
)
conn.execute("DROP TABLE chat_attachment_tombstones_legacy")
conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
except Exception:
conn.execute("ROLLBACK TO SAVEPOINT migrate_chat_attachment_tombstones")
conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
raise
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_attachment_inventory (
message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE,
attachment_id TEXT NOT NULL,
name TEXT NOT NULL,
type TEXT,
content_type TEXT,
size_bytes INTEGER,
PRIMARY KEY(message_id, attachment_id)
) WITHOUT ROWID
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_attachment_inventory_state (
singleton INTEGER NOT NULL PRIMARY KEY CHECK(singleton = 1),
inventory_version INTEGER NOT NULL DEFAULT 0,
dirty INTEGER NOT NULL DEFAULT 1,
backfilled_at INTEGER NOT NULL
)
"""
)
inventory_state_columns = {
row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)")
}
if "inventory_version" not in inventory_state_columns:
conn.execute(
"ALTER TABLE chat_attachment_inventory_state "
"ADD COLUMN inventory_version INTEGER NOT NULL DEFAULT 0"
)
if "dirty" not in inventory_state_columns:
conn.execute(
"ALTER TABLE chat_attachment_inventory_state "
"ADD COLUMN dirty INTEGER NOT NULL DEFAULT 1"
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_insert
AFTER INSERT ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_update
AFTER UPDATE ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_delete
AFTER DELETE ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)"
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_project_id ON chat_threads(project_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_messages_thread_id_created_at ON chat_messages(thread_id, created_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_settings (
key TEXT NOT NULL PRIMARY KEY,
value_json TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_settings (
key TEXT NOT NULL PRIMARY KEY,
value_json TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_settings_quarantine (
id INTEGER PRIMARY KEY AUTOINCREMENT,
key TEXT NOT NULL,
value_json TEXT NOT NULL,
reason TEXT NOT NULL,
quarantined_at TEXT NOT NULL
)
"""
)
# Import ledger inside studio.db (vs a localStorage boolean) so a db wipe re-triggers the
# legacy Dexie import instead of silently hiding threads. Keyed by legacy thread id.
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_legacy_imports (
legacy_thread_id TEXT NOT NULL PRIMARY KEY,
imported_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS prompt_entries (
id TEXT NOT NULL PRIMARY KEY,
name TEXT NOT NULL,
text TEXT NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_entries_created_at ON prompt_entries(created_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS prompt_lists (
id TEXT NOT NULL PRIMARY KEY,
name TEXT NOT NULL,
items_json TEXT NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_lists_created_at ON prompt_lists(created_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_runs (
id TEXT NOT NULL PRIMARY KEY,
owner_subject TEXT NOT NULL,
thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
user_message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE,
assistant_message_id TEXT REFERENCES chat_messages(id) ON DELETE SET NULL,
status TEXT NOT NULL CHECK(status IN (
'planning', 'awaiting_approval', 'queued', 'running', 'paused',
'cancelling', 'cancelled', 'completed', 'failed'
)),
plan_json TEXT,
plan_revision INTEGER NOT NULL DEFAULT 0,
plan_hash TEXT,
config_json TEXT NOT NULL,
cancel_requested INTEGER NOT NULL DEFAULT 0,
lease_owner TEXT,
lease_expires_at INTEGER,
heartbeat_at INTEGER,
retry_count INTEGER NOT NULL DEFAULT 0,
error_message TEXT,
report_text TEXT,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
started_at INTEGER,
completed_at INTEGER,
next_event_seq INTEGER NOT NULL DEFAULT 1
)
"""
)
research_run_cols = {
row[1] for row in conn.execute("PRAGMA table_info(research_runs)").fetchall()
}
if "report_text" not in research_run_cols:
conn.execute("ALTER TABLE research_runs ADD COLUMN report_text TEXT")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_thread_claims (
owner_subject TEXT NOT NULL,
thread_id TEXT NOT NULL PRIMARY KEY REFERENCES chat_threads(id) ON DELETE CASCADE,
created_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
claim_pk = [
row[1]
for row in sorted(
conn.execute("PRAGMA table_info(research_thread_claims)").fetchall(),
key = lambda row: int(row[5] or 0),
)
if int(row[5] or 0) > 0
]
if claim_pk != ["thread_id"]:
# Rebuild the claims table (legacy owner_subject+thread_id PK -> thread_id PK) atomically: in
# autocommit an interruption after CREATE orphaned the rows in _legacy and never re-triggered.
conn.commit()
conn.execute("BEGIN IMMEDIATE")
try:
conn.execute(
"ALTER TABLE research_thread_claims RENAME TO research_thread_claims_legacy"
)
conn.execute(
"""
CREATE TABLE research_thread_claims (
owner_subject TEXT NOT NULL,
thread_id TEXT NOT NULL PRIMARY KEY REFERENCES chat_threads(id) ON DELETE CASCADE,
created_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
conn.execute(
"""INSERT OR IGNORE INTO research_thread_claims
(owner_subject, thread_id, created_at)
SELECT owner_subject, thread_id, created_at
FROM research_thread_claims_legacy
ORDER BY created_at, owner_subject"""
)
conn.execute("DROP TABLE research_thread_claims_legacy")
conn.commit()
except Exception:
conn.rollback()
raise
conn.execute(
"""INSERT OR IGNORE INTO research_thread_claims
(owner_subject, thread_id, created_at)
SELECT owner_subject, thread_id, created_at
FROM research_runs ORDER BY created_at, id"""
)
conn.execute(
"""UPDATE research_runs
SET status='failed', error_message='Superseded by the global thread research claim',
lease_owner=NULL, lease_expires_at=NULL, completed_at=COALESCE(completed_at, updated_at)
WHERE status IN ('planning','awaiting_approval','queued','running','paused','cancelling')
AND EXISTS (
SELECT 1 FROM research_thread_claims c
WHERE c.thread_id=research_runs.thread_id
AND c.owner_subject<>research_runs.owner_subject
)"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_plan_steps (
run_id TEXT NOT NULL REFERENCES research_runs(id) ON DELETE CASCADE,
position INTEGER NOT NULL,
title TEXT NOT NULL,
query TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
result_json TEXT,
started_at INTEGER,
completed_at INTEGER,
PRIMARY KEY(run_id, position)
) WITHOUT ROWID
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_sources (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL REFERENCES research_runs(id) ON DELETE CASCADE,
step_position INTEGER,
url TEXT NOT NULL,
title TEXT,
snippet TEXT,
fetched_at INTEGER NOT NULL,
UNIQUE(run_id, url)
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_document_sources (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL REFERENCES research_runs(id) ON DELETE CASCADE,
step_position INTEGER,
source_key TEXT NOT NULL,
document_id TEXT,
chunk_id TEXT,
filename TEXT NOT NULL,
page INTEGER,
score REAL,
snippet TEXT,
fetched_at INTEGER NOT NULL,
UNIQUE(run_id, source_key)
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS research_events (
run_id TEXT NOT NULL REFERENCES research_runs(id) ON DELETE CASCADE,
seq INTEGER NOT NULL,
event_type TEXT NOT NULL,
data_json TEXT NOT NULL,
created_at INTEGER NOT NULL,
PRIMARY KEY(run_id, seq)
) WITHOUT ROWID
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_research_runs_owner_thread_status "
"ON research_runs(owner_subject, thread_id, status)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_research_runs_lease "
"ON research_runs(status, lease_expires_at)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_research_sources_run ON research_sources(run_id, id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_research_document_sources_run "
"ON research_document_sources(run_id, id)"
)
inventory_state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
# Positional read: works for raw tuple or sqlite3.Row (no row_factory precondition).
if (
inventory_state is None
or inventory_state[0] != _CHAT_ATTACHMENT_INVENTORY_VERSION
or inventory_state[1]
):
_rebuild_chat_attachment_inventory(conn)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
def _prompt_entry_from_row(row: sqlite3.Row) -> dict:
return {
"id": row["id"],
"name": row["name"],
"text": row["text"],
"createdAt": row["created_at"],
"updatedAt": row["updated_at"],
}
def list_prompt_entries() -> list[dict]:
conn = get_connection()
try:
rows = conn.execute("SELECT * FROM prompt_entries ORDER BY created_at DESC").fetchall()
return [_prompt_entry_from_row(r) for r in rows]
finally:
conn.close()
def upsert_prompt_entry(entry: dict) -> dict:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO prompt_entries (id, name, text, created_at, updated_at)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
text = excluded.text,
updated_at = excluded.updated_at
""",
(
entry["id"],
entry["name"],
entry["text"],
entry["createdAt"],
entry["updatedAt"],
),
)
conn.commit()
return entry
finally:
conn.close()
def delete_prompt_entry(entry_id: str) -> None:
conn = get_connection()
try:
conn.execute("DELETE FROM prompt_entries WHERE id = ?", (entry_id,))
conn.commit()
finally:
conn.close()
def bulk_upsert_prompt_entries(entries: list[dict]) -> int:
if not entries:
return 0
conn = get_connection()
try:
conn.executemany(
"""
INSERT INTO prompt_entries (id, name, text, created_at, updated_at)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
text = excluded.text,
updated_at = excluded.updated_at
""",
[(e["id"], e["name"], e["text"], e["createdAt"], e["updatedAt"]) for e in entries],
)
conn.commit()
return len(entries)
finally:
conn.close()
def _prompt_list_from_row(row: sqlite3.Row) -> dict:
return {
"id": row["id"],
"name": row["name"],
"items": json.loads(row["items_json"]),
"createdAt": row["created_at"],
"updatedAt": row["updated_at"],
}
def list_prompt_lists_db() -> list[dict]:
conn = get_connection()
try:
rows = conn.execute("SELECT * FROM prompt_lists ORDER BY created_at DESC").fetchall()
return [_prompt_list_from_row(r) for r in rows]
finally:
conn.close()
def upsert_prompt_list(lst: dict) -> dict:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO prompt_lists (id, name, items_json, created_at, updated_at)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
items_json = excluded.items_json,
updated_at = excluded.updated_at
""",
(
lst["id"],
lst["name"],
json.dumps(lst["items"]),
lst["createdAt"],
lst["updatedAt"],
),
)
conn.commit()
return lst
finally:
conn.close()
def delete_prompt_list_db(list_id: str) -> None:
conn = get_connection()
try:
conn.execute("DELETE FROM prompt_lists WHERE id = ?", (list_id,))
conn.commit()
finally:
conn.close()
def bulk_upsert_prompt_lists(lists: list[dict]) -> int:
if not lists:
return 0
conn = get_connection()
try:
conn.executemany(
"""
INSERT INTO prompt_lists (id, name, items_json, created_at, updated_at)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
items_json = excluded.items_json,
updated_at = excluded.updated_at
""",
[
(
lst["id"],
lst["name"],
json.dumps(lst["items"]),
lst["createdAt"],
lst["updatedAt"],
)
for lst in lists
],
)
conn.commit()
return len(lists)
finally:
conn.close()
# sqlite's own default, kept for every caller that may run on the event loop
_BUSY_TIMEOUT_SECONDS = 5.0
# only for the chat history transactions that run in Starlette's threadpool, where a large clear
# can hold the writer lock past the default and the failure escapes as an unmapped 500
_CONTENDED_BUSY_TIMEOUT_SECONDS = 30.0
def get_connection(busy_timeout_seconds: float = _BUSY_TIMEOUT_SECONDS) -> sqlite3.Connection:
"""Open studio.db with WAL mode, create tables once per process, enable foreign keys."""
global _schema_ready
db_path = studio_db_path()
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path), timeout = busy_timeout_seconds)
conn.row_factory = sqlite3.Row
# foreign_keys is session-scoped; set per connection
conn.execute("PRAGMA foreign_keys=ON")
if not _schema_ready:
with _schema_lock:
if not _schema_ready:
try:
_ensure_schema(conn)
conn.commit()
_schema_ready = True
except Exception:
conn.close()
raise
return conn
def create_run(
id: str,
model_name: str,
dataset_name: str,
config_json: str,
started_at: str,
total_steps: Optional[int],
*,
output_dir: Optional[str] = None,
cancel_requested: bool = False,
resumed_from_run_id: Optional[str] = None,
) -> None:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO training_runs (
id, model_name, dataset_name, config_json, started_at, total_steps,
output_dir, resume_blocked, resumed_from_run_id
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
id,
model_name,
dataset_name,
config_json,
started_at,
total_steps,
None if cancel_requested else output_dir,
int(cancel_requested),
resumed_from_run_id,
),
)
if resumed_from_run_id:
claimed = conn.execute(
"""
UPDATE training_runs SET resume_blocked = 1
WHERE id = ? AND status IN ('stopped', 'error')
AND output_dir = ? AND resume_blocked = 0
""",
(resumed_from_run_id, output_dir),
)
if claimed.rowcount != 1:
raise RuntimeError("Resume source is no longer available")
conn.commit()
finally:
conn.close()
def update_run_total_steps(id: str, total_steps: int) -> None:
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET total_steps = ? WHERE id = ?",
(total_steps, id),
)
conn.commit()
finally:
conn.close()
def update_run_progress(
id: str, step: int, loss: Optional[float], duration_seconds: Optional[float]
) -> None:
"""Update current progress on a running training run (called on each metric flush)."""
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET final_step = ?, final_loss = ?, duration_seconds = ? WHERE id = ?",
(step, loss, duration_seconds, id),
)
conn.commit()
finally:
conn.close()
def finish_run(
id: str,
status: str,
ended_at: str,
final_step: Optional[int],
final_loss: Optional[float],
duration_seconds: Optional[float],
loss_sparkline: Optional[str] = None,
output_dir: Optional[str] = None,
error_message: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
config_json: Optional[str] = None,
) -> None:
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs
SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
duration_seconds = ?, loss_sparkline = ?,
config_json = COALESCE(?, config_json),
output_dir = CASE
WHEN resume_blocked = 1 OR ? = 1 THEN NULL
WHEN ? IS NOT NULL THEN ?
WHEN ? IN ('error', 'stopped') THEN output_dir
ELSE NULL
END,
error_message = ?,
resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END
WHERE id = ? AND status = 'running'
""",
(
status,
ended_at,
final_step,
final_loss,
duration_seconds,
loss_sparkline,
config_json,
int(clear_output_dir),
output_dir,
output_dir,
status,
error_message,
int(clear_output_dir),
int(resume_blocked),
id,
),
)
conn.commit()
finally:
conn.close()
def insert_metrics_batch(run_id: str, metrics: list[dict]) -> None:
if not metrics:
return
conn = get_connection()
try:
conn.executemany(
"""
INSERT INTO training_metrics
(run_id, step, loss, learning_rate, grad_norm, eval_loss, epoch, num_tokens, elapsed_seconds)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(run_id, step) DO UPDATE SET
loss = COALESCE(excluded.loss, loss),
learning_rate = COALESCE(excluded.learning_rate, learning_rate),
grad_norm = COALESCE(excluded.grad_norm, grad_norm),
eval_loss = COALESCE(excluded.eval_loss, eval_loss),
epoch = COALESCE(excluded.epoch, epoch),
num_tokens = COALESCE(excluded.num_tokens, num_tokens),
elapsed_seconds = COALESCE(excluded.elapsed_seconds, elapsed_seconds)
""",
[
(
run_id,
m.get("step"),
m.get("loss"),
m.get("learning_rate"),
m.get("grad_norm"),
m.get("eval_loss"),
m.get("epoch"),
m.get("num_tokens"),
m.get("elapsed_seconds"),
)
for m in metrics
],
)
conn.commit()
finally:
conn.close()
def update_run_display_name(id: str, display_name: Optional[str]) -> None:
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET display_name = ? WHERE id = ?",
(display_name, id),
)
conn.commit()
finally:
conn.close()
def update_run_output_dir(id: str, output_dir: Optional[str]) -> None:
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs SET output_dir = ?
WHERE id = ? AND status = 'running' AND resume_blocked = 0
""",
(output_dir, id),
)
conn.commit()
finally:
conn.close()
def update_run_config_json(id: str, config_json: str) -> bool:
conn = get_connection()
try:
cursor = conn.execute(
"""
UPDATE training_runs SET config_json = ?
WHERE id = ? AND status = 'running'
""",
(config_json, id),
)
conn.commit()
return cursor.rowcount == 1
finally:
conn.close()
def mark_run_cancel_requested(id: str) -> bool:
"""Clear resume/export state only while the exact run is still active."""
conn = get_connection()
try:
cursor = conn.execute(
"""
UPDATE training_runs SET output_dir = NULL, resume_blocked = 1
WHERE id = ? AND status = 'running'
""",
(id,),
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def list_runs(limit: int = 50, offset: int = 0) -> dict:
conn = get_connection()
try:
total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0]
rows = conn.execute(
"""
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked,
CASE
WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
ORDER BY started_at DESC
LIMIT ? OFFSET ?
""",
(limit, offset),
).fetchall()
runs = []
for row in rows:
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug("Failed to parse loss_sparkline for run %s", run.get("id"))
run["loss_sparkline"] = None
runs.append(run)
return {"runs": runs, "total": total}
finally:
conn.close()
def get_run(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT r.*,
CASE
WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
WHERE r.id = ?
""",
(id,),
).fetchone()
if row is None:
return None
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug("Failed to parse loss_sparkline for run %s", id)
run["loss_sparkline"] = None
return run
finally:
conn.close()
def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT r.*,
0 AS resumed_later
FROM training_runs r
WHERE r.output_dir = ?
AND r.status IN ('stopped', 'error')
AND NOT EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
ORDER BY r.started_at DESC
LIMIT 1
""",
(output_dir,),
).fetchone()
if row is None:
return None
run = dict(row)
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug("Failed to parse loss_sparkline for output_dir %s", output_dir)
run["loss_sparkline"] = None
return run
finally:
conn.close()
def get_run_metrics(id: str) -> dict:
"""Return metric arrays for a run, using paired step arrays per metric."""
conn = get_connection()
try:
rows = conn.execute(
"""
SELECT step, loss, learning_rate, grad_norm, eval_loss, epoch,
num_tokens, elapsed_seconds
FROM training_metrics
WHERE run_id = ?
ORDER BY step
""",
(id,),
).fetchall()
step_history: list[int] = []
loss_history: list[float] = []
loss_step_history: list[int] = []
lr_history: list[float] = []
lr_step_history: list[int] = []
grad_norm_history: list[float] = []
grad_norm_step_history: list[int] = []
eval_loss_history: list[float] = []
eval_step_history: list[int] = []
final_epoch: float | None = None
final_num_tokens: int | None = None
for row in rows:
step = row["step"]
step_history.append(step)
if step > 0 and row["loss"] is not None:
loss_history.append(row["loss"])
loss_step_history.append(step)
if step > 0 and row["learning_rate"] is not None:
lr_history.append(row["learning_rate"])
lr_step_history.append(step)
if step > 0 and row["grad_norm"] is not None:
grad_norm_history.append(row["grad_norm"])
grad_norm_step_history.append(step)
if step > 0 and row["eval_loss"] is not None:
eval_loss_history.append(row["eval_loss"])
eval_step_history.append(step)
if row["epoch"] is not None:
final_epoch = row["epoch"]
if row["num_tokens"] is not None:
final_num_tokens = row["num_tokens"]
return {
"step_history": step_history,
"loss_history": loss_history,
"loss_step_history": loss_step_history,
"lr_history": lr_history,
"lr_step_history": lr_step_history,
"grad_norm_history": grad_norm_history,
"grad_norm_step_history": grad_norm_step_history,
"eval_loss_history": eval_loss_history,
"eval_step_history": eval_step_history,
"final_epoch": final_epoch,
"final_num_tokens": final_num_tokens,
}
finally:
conn.close()
def delete_run(id: str) -> None:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
source = conn.execute(
"SELECT resumed_from_run_id FROM training_runs WHERE id = ?",
(id,),
).fetchone()
if source is not None:
# Keep an explicit resume chain connected when its middle row is removed: the surviving tail
# still carries the source's cumulative counters, so profile totals must trace it.
conn.execute(
"""
UPDATE training_runs
SET resumed_from_run_id = ?
WHERE resumed_from_run_id = ?
""",
(source["resumed_from_run_id"], id),
)
conn.execute("DELETE FROM training_runs WHERE id = ?", (id,))
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_other_run_output_dirs(exclude_id: str) -> list[str]:
conn = get_connection()
try:
rows = conn.execute(
"SELECT output_dir FROM training_runs WHERE output_dir IS NOT NULL AND id != ?",
(exclude_id,),
).fetchall()
return [str(row[0]) for row in rows]
finally:
conn.close()
def cleanup_orphaned_runs() -> None:
"""Mark any 'running' rows as errored on startup (server restarted mid-training)."""
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs
SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END,
error_message = CASE
WHEN resume_blocked = 1 THEN NULL
ELSE 'Server restarted during training'
END,
output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END,
ended_at = ?
WHERE status = 'running'
""",
(datetime.now(timezone.utc).isoformat(),),
)
conn.commit()
finally:
conn.close()
def list_scan_folders() -> list[dict]:
conn = get_connection()
try:
rows = conn.execute(
"SELECT id, path, created_at FROM scan_folders ORDER BY created_at"
).fetchall()
return [dict(row) for row in rows]
finally:
conn.close()
def add_scan_folder_with_status(path: str) -> tuple[dict, bool]:
"""Add a custom scan folder and return its row plus whether it was inserted."""
if not path or not path.strip():
raise ValueError("Path cannot be empty")
normalized = os.path.realpath(os.path.expanduser(path.strip()))
# Validate the path is an existing, readable directory before persisting.
if not os.path.exists(normalized):
raise ValueError("Path does not exist")
if not os.path.isdir(normalized):
raise ValueError("Path must be a directory, not a file")
# Reject a local filesystem root ("/", or a bare "C:\\"): registering one seeds the browse
# allowlist above denied system dirs. A UNC share root has none under it, so it stays
# allowed (it was registerable before this guard). Mirrors scan_folders.py.
if is_local_filesystem_root(normalized):
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")
# Windows: normcase for the denylist check but store original casing (e.g. C:\Models).
is_win = platform.system() == "Windows"
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
raise ValueError(f"Path under {prefix} is not allowed")
# Last, so a denied path is never opened: os.access alone passes on folders
# macOS TCC or a Windows ACL still refuses at scan time, which is how a
# registered folder ends up looking empty instead of blocked.
if not is_readable_dir(normalized):
raise ValueError("Path is not readable")
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
# Windows: case-insensitive lookup so C:\Models and c:\models dedup.
if is_win:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE",
(normalized,),
).fetchone()
else:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ?",
(normalized,),
).fetchone()
if existing is not None:
return dict(existing), False
inserted = False
try:
conn.execute(
"INSERT INTO scan_folders (path, created_at) VALUES (?, ?)",
(normalized, now),
)
conn.commit()
inserted = True
except sqlite3.IntegrityError:
pass # duplicate; fall through to SELECT
# Same collation as the pre-check to catch concurrent writes (Windows).
fallback_sql = (
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE"
if is_win
else "SELECT id, path, created_at FROM scan_folders WHERE path = ?"
)
row = conn.execute(fallback_sql, (normalized,)).fetchone()
if row is None:
raise ValueError("Folder was concurrently removed")
return dict(row), inserted
finally:
conn.close()
def add_scan_folder(path: str) -> dict:
"""Add a directory to the custom scan folder list. Returns the row."""
row, _ = add_scan_folder_with_status(path)
return row
def remove_scan_folder(id: int) -> bool:
conn = get_connection()
try:
cursor = conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,))
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def _json_loads(value: str | None, fallback):
if value is None:
return fallback
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return fallback
def _chat_thread_from_row(row: sqlite3.Row, include_settings: bool = True) -> dict:
data = dict(row)
thread = {
"id": data["id"],
"title": data["title"],
"modelType": data["model_type"],
"modelId": data.get("model_id") or "",
"pairId": data.get("pair_id") or None,
"projectId": data.get("project_id") or None,
"archived": bool(data["archived"]),
"createdAt": data["created_at"],
"updatedAt": data.get("updated_at")
if data.get("updated_at") is not None
else data["created_at"],
"openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
"anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
"forkedFromThreadId": data.get("forked_from_thread_id"),
"forkedFromMessageId": data.get("forked_from_message_id"),
}
if include_settings:
thread["settings"] = _json_loads(data.get("settings_json"), None)
return thread
def _chat_project_from_row(row: sqlite3.Row) -> dict:
data = dict(row)
root_path = data.get("root_path")
return {
"id": data["id"],
"name": data["name"],
"instructions": data.get("instructions") or "",
"rootPath": root_path or None,
"sandboxPath": os.path.join(root_path, "sandbox") if root_path else None,
"archived": bool(data["archived"]),
"createdAt": data["created_at"],
"updatedAt": data["updated_at"],
}
def _chat_message_from_row(row: sqlite3.Row) -> dict:
data = dict(row)
message = {
"id": data["id"],
"threadId": data["thread_id"],
"parentId": data.get("parent_id"),
"role": data["role"],
"content": _json_loads(data.get("content_json"), []),
"createdAt": data["created_at"],
}
attachments = _json_loads(data.get("attachments_json"), None)
metadata = _json_loads(data.get("metadata_json"), None)
if attachments is not None:
message["attachments"] = attachments
if metadata is not None:
message["metadata"] = metadata
return message
class ChatThreadDeletedError(RuntimeError):
"""Raised when a stale writer tries to recreate a deleted thread id."""
def _raise_if_chat_thread_deleted(conn: sqlite3.Connection, thread_id: str) -> None:
row = conn.execute(
"SELECT 1 FROM chat_thread_tombstones WHERE id = ?",
(thread_id,),
).fetchone()
if row is not None:
raise ChatThreadDeletedError(thread_id)
def upsert_chat_thread(thread: dict) -> dict:
conn = get_connection(_CONTENDED_BUSY_TIMEOUT_SECONDS)
try:
conn.execute("BEGIN IMMEDIATE")
_raise_if_chat_thread_deleted(conn, thread["id"])
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id, settings_json)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
model_type = excluded.model_type,
model_id = excluded.model_id,
pair_id = excluded.pair_id,
project_id = excluded.project_id,
archived = excluded.archived,
created_at = excluded.created_at,
updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at),
openai_code_exec_container_id = excluded.openai_code_exec_container_id,
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id,
forked_from_thread_id = excluded.forked_from_thread_id,
forked_from_message_id = excluded.forked_from_message_id,
-- an absent snapshot keeps the stored one: most writers rebuild the record without it.
settings_json = COALESCE(excluded.settings_json, chat_threads.settings_json)
""",
(
thread["id"],
thread.get("title") or "New Chat",
thread["modelType"],
thread.get("modelId") or "",
thread.get("pairId"),
thread.get("projectId"),
1 if thread.get("archived") else 0,
int(thread["createdAt"]),
int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None,
thread.get("openaiCodeExecContainerId"),
thread.get("anthropicCodeExecContainerId"),
thread.get("forkedFromThreadId"),
thread.get("forkedFromMessageId"),
json.dumps(thread["settings"]) if thread.get("settings") is not None else None,
),
)
conn.commit()
return get_chat_thread(thread["id"]) or thread
except Exception:
conn.rollback()
raise
finally:
conn.close()
class ChatThreadPreconditionFailed(Exception):
"""The thread changed between the caller reading it and writing it."""
# Same order the title migration picks its opening message in.
# Watermarks kept per thread, one per tab that has written its settings.
_MAX_SETTINGS_WRITERS = 32
_OPENING_USER_MESSAGE = """(
SELECT id FROM chat_messages
WHERE thread_id = ? AND role = 'user'
ORDER BY created_at ASC, id ASC LIMIT 1
) IS ?"""
def update_chat_thread(
id: str,
patch: dict,
expected_title: Optional[str] = None,
expected_opening_message_id: Optional[str] = None,
settings_write: Optional[dict] = None,
) -> Optional[dict]:
"""Patch a thread. With expected_title, the write only lands while the row
still holds that title, so a concurrent rename wins instead of being lost.
expected_opening_message_id guards a title derived from that message: if it
is gone, the write is rejected rather than expanding deleted text."""
allowed = {
"title": ("title", patch.get("title")),
"modelType": ("model_type", patch.get("modelType")),
"modelId": ("model_id", patch.get("modelId")),
"pairId": ("pair_id", patch.get("pairId")),
"projectId": ("project_id", patch.get("projectId")),
"archived": ("archived", 1 if patch.get("archived") else 0),
"createdAt": ("created_at", patch.get("createdAt")),
"updatedAt": ("updated_at", patch.get("updatedAt")),
"openaiCodeExecContainerId": (
"openai_code_exec_container_id",
patch.get("openaiCodeExecContainerId"),
),
"anthropicCodeExecContainerId": (
"anthropic_code_exec_container_id",
patch.get("anthropicCodeExecContainerId"),
),
"forkedFromThreadId": (
"forked_from_thread_id",
patch.get("forkedFromThreadId"),
),
"forkedFromMessageId": (
"forked_from_message_id",
patch.get("forkedFromMessageId"),
),
"settings": (
"settings_json",
json.dumps(patch["settings"]) if patch.get("settings") is not None else None,
),
}
assignments = []
values = []
for key, (column, value) in allowed.items():
if key in patch:
assignments.append(f"{column} = ?")
values.append(value)
if not assignments and settings_write is None:
return get_chat_thread(id)
conn = get_connection()
try:
# The snapshot rides in the same transaction as the guarded metadata write, or a
# rejected precondition returns 409 with the settings change already committed.
conn.execute("BEGIN IMMEDIATE")
# Guards ride in the WHERE clause, so check and write are one statement.
where = ["id = ?"]
guard: list = [id]
if expected_title is not None:
where.append("title = ?")
guard.append(expected_title)
if expected_opening_message_id is not None:
where.append(_OPENING_USER_MESSAGE)
guard += [id, expected_opening_message_id]
guarded = expected_title is not None or expected_opening_message_id is not None
applied = 1
if assignments:
cursor = conn.execute(
f"UPDATE chat_threads SET {', '.join(assignments)} WHERE {' AND '.join(where)}",
(*values, *guard),
)
applied = cursor.rowcount
if guarded and applied == 0:
conn.rollback()
if conn.execute("SELECT 1 FROM chat_threads WHERE id = ?", (id,)).fetchone() is None:
return None
raise ChatThreadPreconditionFailed(id)
if settings_write is not None:
if _write_chat_thread_settings_in_conn(conn, id, **settings_write) is None:
conn.rollback()
return None
conn.commit()
row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone()
if row is None:
return None
return _chat_thread_from_row(row)
finally:
conn.close()
def _write_chat_thread_settings_in_conn(
conn,
id: str,
*,
replace: Optional[dict] = None,
merge: Optional[dict] = None,
clear: bool = False,
seq: Optional[int] = None,
writer: Optional[str] = None,
keep_unreadable = None,
) -> Optional[bool]:
"""The snapshot write itself, on a connection whose transaction the caller owns.
None when the row is gone. True when it wrote, False when an older write from the
same writer was refused; both leave the caller's transaction usable.
The all-or-nothing the caller wants is about FAILURE: a rejected metadata
precondition must not leave the settings committed, and a missing row must write
nothing. A refusal is not a failure. It means this writer has already landed a newer
snapshot, so the row holds what the writer wanted either way, and rolling the
metadata back with it would drop a rename the client sent in the same PATCH and got
a 200 for.
"""
row = conn.execute(
"SELECT settings_json, settings_seqs FROM chat_threads WHERE id = ?",
(id,),
).fetchone()
if row is None:
return None
seqs = _json_loads(row["settings_seqs"], None) if "settings_seqs" in row.keys() else None
seqs = seqs if isinstance(seqs, dict) else {}
if seq is not None and writer is not None:
seen = seqs.get(writer)
if isinstance(seen, int) and seq <= seen:
# This writer has already had a newer snapshot stored; this one is the
# straggler. Held per writer, so a write from another tab in between does
# not wipe the watermark this comparison depends on.
return False
# Re-inserted, not just assigned, so this writer moves to the end and the map
# stays in least-recently-used order.
seqs.pop(writer, None)
seqs[writer] = seq
# One entry per tab that has ever written, so bound it. Evicted by last use and
# never by counter: every session starts its own counter at 1, so comparing them
# across writers would throw out the newest tab and keep long-dead ones, leaving
# the active writer with no watermark for its own stragglers to be refused by.
while len(seqs) > _MAX_SETTINGS_WRITERS:
seqs.pop(next(iter(seqs)))
stored = _json_loads(row["settings_json"], None)
stored = stored if isinstance(stored, dict) else {}
if clear:
settings_json = None
else:
if merge is not None:
base = stored
changes = merge
else:
base = keep_unreadable(stored) if keep_unreadable else {}
changes = replace or {}
settings_json = json.dumps({**base, **changes})
conn.execute(
"UPDATE chat_threads SET settings_json = ?, settings_seqs = ? WHERE id = ?",
(settings_json, json.dumps(seqs) if seqs else None, id),
)
return True
def write_chat_thread_settings(
id: str,
*,
replace: Optional[dict] = None,
merge: Optional[dict] = None,
clear: bool = False,
seq: Optional[int] = None,
writer: Optional[str] = None,
keep_unreadable = None,
) -> Optional[dict]:
"""Write a thread's settings snapshot, reading and merging in one transaction.
Doing the read in the route and the write here lets two requests on the same thread,
two tabs or a tab closing behind an open one, both turn a partial patch into a full
replacement built from the same stale snapshot, and the second one lands on top. The
read, the merge and the write have to be one transaction, so they are.
`writer` and `seq` order the writes, and only ever against the same writer's own
earlier ones: a write is dropped when it comes from the writer whose snapshot is
already stored and carries a seq no newer than it. Two browsers are never compared,
because their clocks and counters have nothing to do with each other and the one that
happened to be behind would have every edit silently refused. Within one writer the
ordering is real, which is the case that needs it: an aborted fetch does not stop a
handler the server has already started.
`keep_unreadable(stored) -> dict` names the part of the stored snapshot the caller
could not read, which a replacement carries forward rather than deleting. Passed in
rather than imported so this module stays free of the wire models.
"""
conn = get_connection()
try:
# IMMEDIATE takes the write lock up front, so the read below cannot be overtaken.
conn.execute("BEGIN IMMEDIATE")
applied = _write_chat_thread_settings_in_conn(
conn,
id,
replace = replace,
merge = merge,
clear = clear,
seq = seq,
writer = writer,
keep_unreadable = keep_unreadable,
)
if applied is None:
conn.rollback()
return None
conn.commit()
row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone()
return _chat_thread_from_row(row) if row is not None else None
finally:
conn.close()
def get_chat_thread(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone()
return _chat_thread_from_row(row) if row is not None else None
finally:
conn.close()
def list_chat_threads(
model_type: str | None = None,
pair_id: str | None = None,
project_id: str | None = None,
include_archived: bool = True,
) -> list[dict]:
clauses = []
values: list[object] = []
if model_type is not None:
clauses.append("model_type = ?")
values.append(model_type)
if pair_id is not None:
clauses.append("pair_id = ?")
values.append(pair_id)
if project_id is not None:
clauses.append("project_id = ?")
values.append(project_id)
if not include_archived:
clauses.append("archived = 0")
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
conn = get_connection()
try:
rows = conn.execute(
f"SELECT * FROM chat_threads {where} "
"ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC",
values,
).fetchall()
# the snapshot is only read when a thread is opened, so it is left out of the listing.
return [_chat_thread_from_row(row, include_settings = False) for row in rows]
finally:
conn.close()
def build_chat_history_export() -> tuple[list[dict], list[dict], list[dict]]:
"""Read projects, threads, and messages from one SQLite snapshot."""
conn = get_connection()
try:
conn.execute("BEGIN")
thread_rows = conn.execute(
"""
SELECT * FROM chat_threads
ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC
"""
).fetchall()
project_rows = conn.execute(
"SELECT * FROM chat_projects ORDER BY updated_at DESC"
).fetchall()
message_rows = conn.execute(
"SELECT * FROM chat_messages ORDER BY created_at ASC, id ASC"
).fetchall()
conn.commit()
return (
[_chat_project_from_row(row) for row in project_rows],
[_chat_thread_from_row(row) for row in thread_rows],
[_chat_message_from_row(row) for row in message_rows],
)
except Exception:
conn.rollback()
raise
finally:
conn.close()
def _reparent_surviving_forks(conn: sqlite3.Connection, deleted_ids: set[str]) -> None:
"""Keep fork lineage connected across any thread deletion path."""
if not deleted_ids:
return
rows = conn.execute("SELECT id, forked_from_thread_id FROM chat_threads").fetchall()
sources = {row["id"]: row["forked_from_thread_id"] for row in rows}
for row in rows:
if row["id"] in deleted_ids or row["forked_from_thread_id"] not in deleted_ids:
continue
source_id = row["forked_from_thread_id"]
seen: set[str] = set()
while source_id in deleted_ids and source_id not in seen:
seen.add(source_id)
parent_id = sources.get(source_id)
if parent_id is None:
# Preserve the deleted root id: sibling forks use it to elect one surviving copy of shared history.
break
source_id = parent_id
conn.execute(
"UPDATE chat_threads SET forked_from_thread_id = ? WHERE id = ?",
(source_id, row["id"]),
)
def _tombstone_chat_threads(conn: sqlite3.Connection, thread_ids: Iterable[str]) -> None:
deleted_at = int(datetime.now(timezone.utc).timestamp() * 1000)
conn.executemany(
"""
INSERT INTO chat_thread_tombstones (id, deleted_at)
VALUES (?, ?)
ON CONFLICT(id) DO UPDATE SET deleted_at = excluded.deleted_at
""",
[(thread_id, deleted_at) for thread_id in sorted(set(thread_ids))],
)
# Only these states can already own a worker. Queued, paused, and approval rows are removed before
# they can be claimed, so signalling them would only leave unused supervisor cancellation events.
_ACTIVE_RESEARCH_RUN_STATUSES = (
"planning",
"running",
"cancelling",
)
def _active_research_run_ids(
conn: sqlite3.Connection, thread_ids: set[str] | None = None
) -> list[str]:
# lease_owner is the other half of "can already own a worker": a planning run that no worker
# has claimed yet never will once its row is gone, and signalling it would leave a
# cancellation event in the supervisor that nothing is left to consume.
status_placeholders = ",".join("?" for _ in _ACTIVE_RESEARCH_RUN_STATUSES)
if thread_ids is None:
rows = conn.execute(
f"""
SELECT id, created_at FROM research_runs
WHERE status IN ({status_placeholders}) AND lease_owner IS NOT NULL
""",
_ACTIVE_RESEARCH_RUN_STATUSES,
).fetchall()
else:
rows = []
sorted_thread_ids = sorted(thread_ids)
for start in range(0, len(sorted_thread_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = sorted_thread_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
thread_placeholders = ",".join("?" for _ in chunk)
rows.extend(
conn.execute(
f"SELECT id, created_at FROM research_runs "
f"WHERE thread_id IN ({thread_placeholders}) "
f"AND status IN ({status_placeholders}) "
f"AND lease_owner IS NOT NULL",
(*chunk, *_ACTIVE_RESEARCH_RUN_STATUSES),
).fetchall()
)
return [row["id"] for row in sorted(rows, key = lambda row: (row["created_at"], row["id"]))]
def delete_chat_threads_with_active_research_runs(ids: list[str]) -> list[str]:
if not ids:
return []
conn = get_connection(_CONTENDED_BUSY_TIMEOUT_SECONDS)
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
thread_ids = set(ids)
active_research_run_ids = _active_research_run_ids(conn, thread_ids)
_reparent_surviving_forks(conn, thread_ids)
# Record the delete even when no row exists yet. A late POST carrying the same unique id
# must not recreate a thread after this request has confirmed deletion.
_tombstone_chat_threads(conn, thread_ids)
conn.executemany(
"DELETE FROM chat_attachment_tombstones WHERE thread_id = ?",
[(id,) for id in ids],
)
conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids])
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return active_research_run_ids
except Exception:
conn.rollback()
raise
finally:
conn.close()
def delete_chat_threads(ids: list[str]) -> list[str]:
"""Delete threads and return the active research runs removed with them."""
return delete_chat_threads_with_active_research_runs(ids)
def clear_chat_history_with_active_research_runs(
additional_thread_ids: Iterable[str] = (), operation_id: Optional[str] = None
) -> tuple[list[str], list[str]]:
removed, active_runs = clear_chat_history(
additional_thread_ids,
operation_id = operation_id,
)
return active_runs, removed
def clear_chat_history(
additional_thread_ids: Iterable[str] = (), operation_id: Optional[str] = None
) -> "tuple[list[str], list[str]]":
"""Delete every chat thread. Returns (thread ids removed, research runs cascaded).
Both taken inside the same transaction: another process can add a thread
between a listing and this call, its sandbox has to be cleaned up too, and
after the cascade nothing can tell the supervisor which runs to stop.
"""
conn = get_connection(_CONTENDED_BUSY_TIMEOUT_SECONDS)
try:
conn.execute("BEGIN IMMEDIATE")
if operation_id is not None:
completed = conn.execute(
"""
SELECT deleted_thread_ids_json
FROM chat_clear_operations WHERE id = ?
""",
(operation_id,),
).fetchone()
if completed is not None:
conn.commit()
# The original request already signalled its workers. Replaying that signal can
# leave a cancellation event behind after the worker has exited.
return list(json.loads(completed["deleted_thread_ids_json"])), []
_ensure_chat_attachment_inventory_current(conn)
removed = sorted(str(row[0]) for row in conn.execute("SELECT id FROM chat_threads"))
status_placeholders = ",".join("?" for _ in _ACTIVE_RESEARCH_RUN_STATUSES)
active_runs = [
row["id"]
for row in conn.execute(
f"SELECT id FROM research_runs WHERE status IN ({status_placeholders}) "
"AND lease_owner IS NOT NULL ORDER BY created_at, id",
_ACTIVE_RESEARCH_RUN_STATUSES,
)
]
# Fence pending frontend writes and legacy-only ids in the same transaction as the clear.
_tombstone_chat_threads(conn, sorted(set(additional_thread_ids) | set(removed)))
conn.execute("DELETE FROM chat_attachment_tombstones")
conn.execute("DELETE FROM chat_threads")
_mark_chat_attachment_inventory_clean(conn)
if operation_id is not None:
conn.execute(
"""
INSERT INTO chat_clear_operations (
id, active_research_run_ids_json, deleted_thread_ids_json, cleared_at
) VALUES (?, ?, ?, ?)
""",
(
operation_id,
json.dumps(active_runs),
json.dumps(removed),
int(datetime.now(timezone.utc).timestamp() * 1000),
),
)
conn.commit()
return removed, active_runs
except Exception:
conn.rollback()
raise
finally:
conn.close()
def count_chat_threads() -> int:
conn = get_connection()
try:
return int(conn.execute("SELECT COUNT(*) FROM chat_threads").fetchone()[0])
finally:
conn.close()
def upsert_chat_project(project: dict) -> dict:
existing = get_chat_project(project["id"])
root_path = existing.get("rootPath") if existing else None
if not root_path:
root_path = _default_project_root(project)
root_path = _ensure_project_workspace(root_path)
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO chat_projects
(id, name, instructions, root_path, archived, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
instructions = excluded.instructions,
root_path = COALESCE(chat_projects.root_path, excluded.root_path),
archived = excluded.archived,
created_at = excluded.created_at,
updated_at = excluded.updated_at
""",
(
project["id"],
project["name"],
project.get("instructions") or "",
root_path,
1 if project.get("archived") else 0,
int(project["createdAt"]),
int(project["updatedAt"]),
),
)
conn.commit()
return get_chat_project(project["id"]) or project
finally:
conn.close()
def update_chat_project(id: str, patch: dict) -> Optional[dict]:
allowed = {
"name": ("name", patch.get("name")),
"instructions": ("instructions", patch.get("instructions")),
"archived": ("archived", 1 if patch.get("archived") else 0),
"createdAt": ("created_at", patch.get("createdAt")),
"updatedAt": ("updated_at", patch.get("updatedAt")),
}
assignments = []
values = []
for key, (column, value) in allowed.items():
if key in patch:
assignments.append(f"{column} = ?")
values.append(value)
if not assignments:
return get_chat_project(id)
conn = get_connection()
try:
conn.execute(
f"UPDATE chat_projects SET {', '.join(assignments)} WHERE id = ?",
(*values, id),
)
conn.commit()
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
return _chat_project_from_row(row) if row is not None else None
finally:
conn.close()
def ensure_chat_project_workspace(id: str) -> Optional[dict]:
project = get_chat_project(id)
if project is None:
return None
root_path = project.get("rootPath") or _default_project_root(project)
root_path = _ensure_project_workspace(root_path)
# a delete running in another threadpool worker can drop the row at any point before the
# directory is created, so confirm the project outlived the create rather than trusting a
# pre-create snapshot. Removing the directory here is not this function's call: only the
# delete path knows whether the user asked to keep the files, and the row may have had a
# populated workspace already. An empty directory left behind is the cheaper outcome.
project = get_chat_project(id)
if project is None:
return None
if project.get("rootPath") == root_path:
return project
conn = get_connection()
try:
conn.execute(
"UPDATE chat_projects SET root_path = ? WHERE id = ?",
(root_path, id),
)
conn.commit()
finally:
conn.close()
return get_chat_project(id)
def get_chat_project(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
return _chat_project_from_row(row) if row is not None else None
finally:
conn.close()
def list_chat_projects(include_archived: bool = False) -> list[dict]:
conn = get_connection()
try:
where = "" if include_archived else "WHERE archived = 0"
rows = conn.execute(
f"SELECT * FROM chat_projects {where} ORDER BY updated_at DESC"
).fetchall()
return [_chat_project_from_row(row) for row in rows]
finally:
conn.close()
def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
conn = get_connection(_CONTENDED_BUSY_TIMEOUT_SECONDS)
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
if row is None:
conn.rollback()
return None
project = _chat_project_from_row(row)
thread_ids = {
thread["id"]
for thread in conn.execute(
"SELECT id FROM chat_threads WHERE project_id = ?",
(id,),
)
}
# Read before the cascade removes them: afterwards nothing can tell the
# supervisor which runs to stop, and a worker keeps doing model, web and
# RAG work for a project that is gone.
active_runs = (
[
row["id"]
for row in conn.execute(
"SELECT id FROM research_runs WHERE thread_id IN ({}) "
"AND status IN ({}) AND lease_owner IS NOT NULL "
"ORDER BY created_at, id".format(
",".join("?" for _ in thread_ids) or "NULL",
",".join("?" for _ in _ACTIVE_RESEARCH_RUN_STATUSES),
),
(*tuple(sorted(thread_ids)), *_ACTIVE_RESEARCH_RUN_STATUSES),
)
]
if thread_ids
else []
)
_reparent_surviving_forks(conn, thread_ids)
# Fence the exact membership selected by this transaction so a late writer cannot
# recreate a project member after the project and its workspace are gone.
_tombstone_chat_threads(conn, thread_ids)
conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,))
conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,))
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
if delete_files:
_delete_project_workspace(project)
# The membership this transaction actually deleted, which is not the
# caller's earlier listing when a chat was moved in between the two.
project = dict(project)
project["memberIds"] = sorted(thread_ids)
project["activeResearchRunIds"] = active_runs
return project
except Exception:
conn.rollback()
raise
finally:
conn.close()
def delete_chat_project_with_active_research_runs(id: str) -> tuple[Optional[dict], list[str]]:
project = delete_chat_project(id, delete_files = False)
if project is None:
return None, []
return project, list(project.get("activeResearchRunIds") or [])
class ChatMessageConflictError(RuntimeError):
"""Raised when a chat message id already belongs to another thread."""
class ChatMessageProtectedError(RuntimeError):
"""Raised when pruning would remove a message owned by a durable feature."""
class CorruptSettingsError(RuntimeError):
"""Raised when a partial settings patch would overwrite corrupt settings."""
def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]:
try:
return True, json.loads(value_json)
except (json.JSONDecodeError, TypeError) as exc:
logger.warning(
"Corrupt chat_settings JSON; quarantining key=%s error=%s",
key,
exc,
)
return False, None
def _load_chat_settings_for_merge(conn: sqlite3.Connection) -> tuple[dict[str, Any], set[str]]:
rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall()
current: dict[str, Any] = {}
corrupt: set[str] = set()
now = datetime.now(timezone.utc).isoformat()
for row in rows:
ok, value = _parse_chat_setting_json(row["key"], row["value_json"])
if ok:
current[row["key"]] = value
continue
corrupt.add(row["key"])
conn.execute(
"""
INSERT INTO chat_settings_quarantine
(key, value_json, reason, quarantined_at)
VALUES (?, ?, ?, ?)
""",
(row["key"], row["value_json"], "json_decode_error", now),
)
conn.execute(
"DELETE FROM chat_settings WHERE key = ? AND value_json = ?",
(row["key"], row["value_json"]),
)
return current, corrupt
def _raise_if_chat_message_thread_conflicts(
conn: sqlite3.Connection, thread_id: str, message_ids: list[str]
) -> None:
unique_ids = list(dict.fromkeys(message_ids))
if not unique_ids:
return
conflicts: list[str] = []
for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT id FROM chat_messages
WHERE id IN ({placeholders}) AND thread_id != ?
ORDER BY id
""",
(*chunk, thread_id),
).fetchall()
conflicts.extend(row["id"] for row in rows)
if conflicts:
preview = ", ".join(conflicts[:5])
suffix = "" if len(conflicts) <= 5 else f" (+{len(conflicts) - 5} more)"
raise ChatMessageConflictError(
f"Message id already belongs to another thread: {preview}{suffix}"
)
def _bump_chat_thread_updated_at(
conn: sqlite3.Connection, thread_id: str, message_created_at: int
) -> None:
conn.execute(
"""
UPDATE chat_threads
SET updated_at = MAX(COALESCE(updated_at, created_at), ?)
WHERE id = ?
""",
(message_created_at, thread_id),
)
def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None:
"""Set updated_at from the remaining messages, floored at created_at.
Unlike the ratchet-only bump, this can lower updated_at -- needed after
pruning, which may delete the thread's newest message.
"""
conn.execute(
"""
UPDATE chat_threads
SET updated_at = MAX(
COALESCE(
(
SELECT MAX(m.created_at) FROM chat_messages m
WHERE m.thread_id = chat_threads.id
),
created_at
),
created_at
)
WHERE id = ?
""",
(thread_id,),
)
def _research_message_ids(conn: sqlite3.Connection, thread_id: str) -> set[str]:
return {
str(message_id)
for row in conn.execute(
"SELECT user_message_id, assistant_message_id FROM research_runs WHERE thread_id = ?",
(thread_id,),
).fetchall()
for message_id in row
if message_id is not None
}
def _surviving_parent_id(
conn: sqlite3.Connection, thread_id: str, message_id: str, pruned: set
) -> "str | None":
"""The stored ancestor a message relinks to once `pruned` is deleted, or None at the root.
Walking the stored chain server side is what makes the relink allowance safe: the expected
parent is derived from rows the server already holds, so a client cannot smuggle an arbitrary
link past the guard by claiming its old parent went away.
"""
seen = {message_id}
row = conn.execute(
"SELECT parent_id FROM chat_messages WHERE thread_id = ? AND id = ?",
(thread_id, message_id),
).fetchone()
parent = row["parent_id"] if row is not None else None
while parent and str(parent) in pruned:
if str(parent) in seen:
# A cycle can only come from a corrupt thread; stop rather than spin.
return None
seen.add(str(parent))
row = conn.execute(
"SELECT parent_id FROM chat_messages WHERE thread_id = ? AND id = ?",
(thread_id, str(parent)),
).fetchone()
parent = row["parent_id"] if row is not None else None
# Normalized the way the caller reads parentId (`or None`), so an empty stored parent_id
# cannot make the two disagree, and a self-link left by a corrupt chain resolves to the root.
survivor = str(parent) if parent else None
return None if survivor == message_id else survivor
def _research_message_would_change(
conn: sqlite3.Connection,
thread_id: str,
message: dict,
pruned: set = frozenset(),
) -> bool:
message_id = str(message["id"])
row = conn.execute(
"SELECT parent_id, role, content_json, metadata_json, attachments_json, created_at "
"FROM chat_messages WHERE thread_id = ? AND id = ?",
(thread_id, message_id),
).fetchone()
if row is None:
return False
def canon(value: object) -> str | None:
return json.dumps(value, sort_keys = True) if value is not None else None
stored_parent = row["parent_id"] or None
sent_parent = message.get("parentId") or None
parent_changed = sent_parent != stored_parent
if parent_changed and stored_parent is not None and str(stored_parent) in pruned:
# The same sync is deleting this message's parent, so the client is repairing a link this
# request is about to break rather than editing a protected message. Only the relink the
# server itself would compute is accepted; anything else is still a rejected edit.
parent_changed = sent_parent != _surviving_parent_id(conn, thread_id, message_id, pruned)
# created_at is compared too: without it a client could re-upsert a protected message with an
# unchanged body but a different timestamp and silently reorder the research prompt/response
# pair. Absent createdAt defaults to the stored value (a no-op re-sync).
return (
canon(message.get("content", [])) != canon(json.loads(row["content_json"] or "[]"))
or canon(message.get("metadata"))
!= canon(json.loads(row["metadata_json"]) if row["metadata_json"] else None)
or canon(message.get("attachments"))
!= canon(json.loads(row["attachments_json"]) if row["attachments_json"] else None)
or parent_changed
or str(message.get("role")) != str(row["role"])
or int(message.get("createdAt", row["created_at"])) != int(row["created_at"])
)
def _guard_research_messages(
conn: sqlite3.Connection,
thread_id: str,
messages: list[dict],
pruned: set = frozenset(),
) -> None:
protected = _research_message_ids(conn, thread_id)
if not protected:
return
for message in messages:
if str(message["id"]) in protected and _research_message_would_change(
conn, thread_id, message, pruned
):
raise ChatMessageProtectedError(
"Research prompts and responses are server-managed and cannot be edited"
)
_CONTENT_PART_ID_PREFIX = "content-part-sha256-"
_URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:")
def _is_locally_stored_blob(value: str) -> bool:
"""True for data URIs or bare base64, never external/blob URI references."""
candidate = value.lstrip()
if not candidate:
return False
if candidate[:5].lower() == "data:":
return True
if candidate.startswith(("//", "\\\\")):
return False
return _URI_SCHEME_RE.match(candidate) is None
def _managed_content_part_payload(part: dict) -> Optional[tuple[str, Any]]:
"""Return the locally stored blob payload used to identify a content part."""
image = part.get("image")
if isinstance(image, str) and image[:5].lower() == "data:":
return "image", image
audio = part.get("audio")
if isinstance(audio, str) and _is_locally_stored_blob(audio):
return "audio", audio
if isinstance(audio, dict):
data = audio.get("data")
if isinstance(data, str) and _is_locally_stored_blob(data):
return "audio", audio
return None
def _content_part_id(part: dict) -> Optional[str]:
"""Stable managed id derived from blob data, without mutating inference content."""
payload = _managed_content_part_payload(part)
if payload is None:
return None
canonical = json.dumps(
payload,
ensure_ascii = False,
separators = (",", ":"),
sort_keys = True,
).encode("utf-8")
return f"{_CONTENT_PART_ID_PREFIX}{hashlib.sha256(canonical).hexdigest()}"
def _chat_attachment_tombstones_for_messages(
conn: sqlite3.Connection, thread_id: str, message_ids: list[str]
) -> dict[str, set[str]]:
tombstones = {message_id: set() for message_id in message_ids}
unique_ids = list(dict.fromkeys(message_ids))
for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT message_id, attachment_id
FROM chat_attachment_tombstones
WHERE thread_id = ? AND message_id IN ({placeholders})
""",
(thread_id, *chunk),
).fetchall()
for row in rows:
tombstones[row["message_id"]].add(row["attachment_id"])
return tombstones
def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict:
"""Strip uploads previously deleted through the Data tab from a stale write."""
if not tombstones:
return message
reconciled = dict(message)
attachments = message.get("attachments")
if isinstance(attachments, list):
reconciled["attachments"] = [
attachment
for attachment in attachments
if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones)
]
content = message.get("content")
if isinstance(content, list):
reconciled["content"] = [
part
for part in content
if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones)
]
return reconciled
def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]:
"""Keep untyped legacy/import metadata safe for SQLite binding."""
if value is None:
return fallback
if isinstance(value, str):
return value or fallback
if isinstance(value, (bool, int, float)):
return str(value)
# Objects and arrays are not useful display metadata and sqlite3 rejects binding them.
return fallback
def _chat_attachment_inventory_entries(
attachments_json: Optional[str],
content_json: Optional[str],
tombstones: Optional[set[str]] = None,
) -> list[dict]:
tombstones = tombstones or set()
attachments = _json_loads(attachments_json, None)
if not isinstance(attachments, list):
attachments = []
attachments = [
attachment
for attachment in attachments
if isinstance(attachment, dict) and attachment.get("id")
]
represented_parts = {
part_id
for attachment in attachments
for part in _attachment_content_parts(attachment)
for part_id in [_content_part_id(part)]
if part_id is not None
}
attachments.extend(
attachment
for attachment in _content_part_attachments(content_json)
if attachment["id"] not in represented_parts
)
entries: list[dict] = []
seen: set[str] = set()
for attachment in attachments:
attachment_id = str(attachment["id"])
if attachment_id in seen or attachment_id in tombstones:
continue
seen.add(attachment_id)
entries.append(
{
"id": attachment_id,
"name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"),
"type": _chat_attachment_metadata_text(attachment.get("type")),
"contentType": _chat_attachment_metadata_text(attachment.get("contentType")),
"sizeBytes": _chat_attachment_size_bytes(attachment),
}
)
return entries
def count_chat_message_attachments(
attachments_json: Optional[str], content_json: Optional[str]
) -> int:
"""Count distinct user-visible uploads represented by a chat message."""
return len(_chat_attachment_inventory_entries(attachments_json, content_json))
def _replace_chat_attachment_inventory(
conn: sqlite3.Connection,
message_id: str,
attachments_json: Optional[str],
content_json: Optional[str],
tombstones: Optional[set[str]] = None,
) -> None:
conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,))
entries = _chat_attachment_inventory_entries(
attachments_json,
content_json,
tombstones,
)
conn.executemany(
"""
INSERT INTO chat_attachment_inventory
(message_id, attachment_id, name, type, content_type, size_bytes)
VALUES (?, ?, ?, ?, ?, ?)
""",
[
(
message_id,
entry["id"],
entry["name"],
entry["type"],
entry["contentType"],
entry["sizeBytes"],
)
for entry in entries
],
)
def _mark_chat_attachment_inventory_clean(conn: sqlite3.Connection) -> None:
conn.execute(
"""
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, ?, 0, ?)
ON CONFLICT(singleton) DO UPDATE SET
inventory_version = excluded.inventory_version,
dirty = 0,
backfilled_at = excluded.backfilled_at
""",
(
_CHAT_ATTACHMENT_INVENTORY_VERSION,
int(datetime.now(timezone.utc).timestamp() * 1000),
),
)
def _rebuild_chat_attachment_inventory(conn: sqlite3.Connection) -> None:
"""Rebuild after schema upgrade or a write from an older Studio build."""
conn.execute("DELETE FROM chat_attachment_inventory")
tombstones: dict[tuple[str, str], set[str]] = {}
for row in conn.execute(
"SELECT thread_id, message_id, attachment_id FROM chat_attachment_tombstones"
).fetchall():
tombstones.setdefault((row["thread_id"], row["message_id"]), set()).add(
row["attachment_id"]
)
rows = conn.execute(
"SELECT id, thread_id, attachments_json, content_json FROM chat_messages"
).fetchall()
for row in rows:
_replace_chat_attachment_inventory(
conn,
row["id"],
row["attachments_json"],
row["content_json"],
tombstones.get((row["thread_id"], row["id"]), set()),
)
def _ensure_chat_attachment_inventory_current(conn: sqlite3.Connection) -> None:
state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
if (
state is not None
and state["inventory_version"] == _CHAT_ATTACHMENT_INVENTORY_VERSION
and not state["dirty"]
):
return
owns_transaction = not conn.in_transaction
if owns_transaction:
conn.execute("BEGIN IMMEDIATE")
try:
state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
if (
state is None
or state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION
or state["dirty"]
):
_rebuild_chat_attachment_inventory(conn)
_mark_chat_attachment_inventory_clean(conn)
if owns_transaction:
conn.commit()
except Exception:
if owns_transaction:
conn.rollback()
raise
def upsert_chat_message(message: dict, *, allow_research_update: bool = False) -> dict:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
if not allow_research_update:
_guard_research_messages(conn, message["threadId"], [message])
_raise_if_chat_message_thread_conflicts(
conn,
message["threadId"],
[message["id"]],
)
tombstones = _chat_attachment_tombstones_for_messages(
conn,
message["threadId"],
[message["id"]],
)
reconciled = _reconcile_chat_message_uploads(
message,
tombstones.get(message["id"], set()),
)
content_json = json.dumps(reconciled.get("content", []))
attachments_json = (
json.dumps(reconciled.get("attachments"))
if reconciled.get("attachments") is not None
else None
)
conn.execute(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
parent_id = excluded.parent_id,
role = excluded.role,
content_json = excluded.content_json,
attachments_json = excluded.attachments_json,
metadata_json = excluded.metadata_json,
created_at = excluded.created_at
WHERE excluded.thread_id = chat_messages.thread_id
""",
(
reconciled["id"],
reconciled["threadId"],
reconciled.get("parentId"),
reconciled["role"],
content_json,
attachments_json,
json.dumps(reconciled.get("metadata"))
if reconciled.get("metadata") is not None
else None,
int(reconciled["createdAt"]),
),
)
_replace_chat_attachment_inventory(
conn,
reconciled["id"],
attachments_json,
content_json,
)
_bump_chat_thread_updated_at(
conn,
reconciled["threadId"],
int(reconciled["createdAt"]),
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return reconciled
except Exception:
conn.rollback()
raise
finally:
conn.close()
def sync_chat_messages(
thread_id: str,
messages: list[dict],
prune_missing: bool = False,
*,
allow_research_update: bool = False,
) -> list[dict]:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
# Research messages are server-managed: keep the server record rather than reject the
# batch on client drift. No _guard_research_messages call here as a result -- these ids
# never reach it. upsert_chat_message still guards, so the single-message route keeps
# rejecting edits.
research_ids = _research_message_ids(conn, thread_id)
# The rows this sync will delete, computed before the upsert so a relink forced by
# that deletion can be told apart from an edit. Research ids are subtracted because
# the delete below exempts them: counting one as pruned would walk the reseat past a
# parent that actually survives, detaching a research turn from its own prompt.
pruned: set = set()
if prune_missing:
retained = {str(m["id"]) for m in messages}
pruned = (
{
str(row["id"])
for row in conn.execute(
"SELECT id FROM chat_messages WHERE thread_id = ?",
(thread_id,),
).fetchall()
}
- retained
- research_ids
)
protected = set() if allow_research_update else research_ids
messages = [m for m in messages if str(m["id"]) not in protected]
# Content is dropped, structure is not: the prune below can delete a research
# message's parent, and a dangling parent makes the whole thread unimportable. The
# replacement is walked from the stored chain, never taken from the client.
#
# Candidates come from research_ids rather than `protected` because the delete exempts
# research rows whatever allow_research_update says, so a narrower set would leave one
# dangling. Ids the batch itself writes are excluded: an authorized caller reparenting
# a research row must not have that overwritten by the repair.
reseat_candidates = research_ids - {str(m["id"]) for m in messages}
reseat_parents = {
message_id: _surviving_parent_id(conn, thread_id, message_id, pruned)
for message_id, stored_parent in _parents_of(conn, thread_id, reseat_candidates).items()
if stored_parent is not None and stored_parent in pruned
}
_raise_if_chat_message_thread_conflicts(
conn,
thread_id,
[m["id"] for m in messages],
)
tombstones = _chat_attachment_tombstones_for_messages(
conn,
thread_id,
[m["id"] for m in messages],
)
reconciled_messages = [
_reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages
]
serialized_messages = [
(
m,
json.dumps(m.get("content", [])),
json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
)
for m in reconciled_messages
]
conn.executemany(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
parent_id = excluded.parent_id,
role = excluded.role,
content_json = excluded.content_json,
attachments_json = excluded.attachments_json,
metadata_json = excluded.metadata_json,
created_at = excluded.created_at
WHERE excluded.thread_id = chat_messages.thread_id
""",
[
(
m["id"],
thread_id,
m.get("parentId"),
m["role"],
content_json,
attachments_json,
json.dumps(m.get("metadata")) if m.get("metadata") is not None else None,
int(m["createdAt"]),
)
for m, content_json, attachments_json in serialized_messages
],
)
for m, content_json, attachments_json in serialized_messages:
_replace_chat_attachment_inventory(
conn,
m["id"],
attachments_json,
content_json,
)
if prune_missing:
retained_ids = {m["id"] for m in reconciled_messages}
existing_ids = {
row["id"]
for row in conn.execute(
"SELECT id FROM chat_messages WHERE thread_id = ?",
(thread_id,),
).fetchall()
}
# Update permission is not delete permission: prune-exempt even for
# allow_research_update callers.
missing_ids = sorted(existing_ids - retained_ids - research_ids)
for start in range(0, len(missing_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = missing_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
conn.execute(
f"DELETE FROM chat_messages WHERE thread_id = ? AND id IN ({placeholders})",
(thread_id, *chunk),
)
_reseat_protected_messages(conn, thread_id, reseat_parents)
_recompute_chat_thread_updated_at(conn, thread_id)
elif reconciled_messages:
_bump_chat_thread_updated_at(
conn,
thread_id,
max(int(m["createdAt"]) for m in reconciled_messages),
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return list_chat_messages(thread_id)
except (ChatMessageConflictError, ChatMessageProtectedError):
conn.rollback()
raise
except sqlite3.Error:
logger.exception("Failed to sync chat messages for thread %s", thread_id)
conn.rollback()
raise
finally:
conn.close()
def _parents_of(conn, thread_id: str, message_ids: set) -> dict:
"""Stored parent of each id in *message_ids*, for rows that exist."""
if not message_ids:
return {}
return {
str(row["id"]): (str(row["parent_id"]) if row["parent_id"] else None)
for row in conn.execute(
"SELECT id, parent_id FROM chat_messages WHERE thread_id = ?",
(thread_id,),
).fetchall()
if str(row["id"]) in message_ids
}
def _reseat_protected_messages(conn, thread_id: str, reseat_parents: dict) -> None:
"""Point protected messages at the ancestor that survived the prune.
Their own rows survive it, the parents they pointed at need not, and a dangling parent
makes the whole thread unimportable on the next load rather than just that turn.
"""
for message_id, parent_id in reseat_parents.items():
conn.execute(
"UPDATE chat_messages SET parent_id = ? WHERE thread_id = ? AND id = ?",
(parent_id, thread_id, message_id),
)
_RESEARCH_LINK_KEYS = {
"researchRunId",
"researchRun",
"researchStatus",
"researchPlanRevision",
"serverManaged",
}
def _detach_research_message_json(
content_json: str, metadata_json: str | None
) -> tuple[str, str | None]:
content = _json_loads(content_json, [])
metadata = _json_loads(metadata_json, None)
custom = metadata.get("custom") if isinstance(metadata, dict) else None
linked = (
isinstance(metadata, dict)
and any(key in metadata for key in _RESEARCH_LINK_KEYS)
or isinstance(custom, dict)
and any(key in custom for key in _RESEARCH_LINK_KEYS)
or isinstance(content, list)
and any(
isinstance(part, dict) and any(key in part for key in _RESEARCH_LINK_KEYS)
for part in content
)
)
if not linked:
return content_json, metadata_json
if isinstance(content, list):
content = [
{key: value for key, value in part.items() if key not in _RESEARCH_LINK_KEYS}
if isinstance(part, dict)
else part
for part in content
]
if isinstance(metadata, dict):
metadata = {key: value for key, value in metadata.items() if key not in _RESEARCH_LINK_KEYS}
custom = metadata.get("custom")
if isinstance(custom, dict):
metadata["custom"] = {
key: value for key, value in custom.items() if key not in _RESEARCH_LINK_KEYS
}
return (
json.dumps(content, ensure_ascii = False),
json.dumps(metadata, ensure_ascii = False) if metadata is not None else None,
)
def fork_chat_thread(
source_thread_id: str,
branch_message_id: str,
new_thread_id: str,
new_title: str,
created_at: int,
id_factory,
) -> Optional[dict]:
"""Atomically clone thread + ancestor msgs `[root..branch_message_id]`
into a new thread. Returns the new thread dict (with messages copied)
or None if source missing.
Reset both code-exec container ids -- per-provider snapshot is handled
by the route layer (best-effort, OpenAI only).
`id_factory()` produces fresh message uuids; injected for testability.
"""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
_raise_if_chat_thread_deleted(conn, new_thread_id)
src = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,)
).fetchone()
if src is None:
conn.rollback()
return None
# Verify branch msg belongs to source thread.
branch_row = conn.execute(
"SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
(source_thread_id, branch_message_id),
).fetchone()
if branch_row is None:
conn.rollback()
return None
# Walk ancestry from branch msg back to root via parent_id chain.
ancestry: list[sqlite3.Row] = []
cursor_row = branch_row
seen: set[str] = set()
while cursor_row is not None and cursor_row["id"] not in seen:
ancestry.append(cursor_row)
seen.add(cursor_row["id"])
parent = cursor_row["parent_id"]
if not parent:
break
cursor_row = conn.execute(
"SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
(source_thread_id, parent),
).fetchone()
ancestry.reverse() # root .. branch msg
# Map old msg id -> new msg id for parent_id rewriting.
id_map: dict[str, str] = {row["id"]: id_factory() for row in ancestry}
src_dict = dict(src)
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at,
openai_code_exec_container_id, anthropic_code_exec_container_id,
forked_from_thread_id, forked_from_message_id, settings_json)
VALUES (?, ?, ?, ?, ?, ?, 0, ?, NULL, NULL, ?, ?, ?)
""",
(
new_thread_id,
new_title,
src_dict["model_type"],
src_dict.get("model_id") or "",
None, # pairId: forks always standalone (compare-mode disabled v1)
src_dict.get("project_id"),
int(created_at),
source_thread_id,
branch_message_id,
src_dict.get("settings_json"),
),
)
fork_messages = []
for row in ancestry:
content_json, metadata_json = _detach_research_message_json(
row["content_json"], row["metadata_json"]
)
fork_messages.append(
(
id_map[row["id"]],
new_thread_id,
id_map.get(row["parent_id"]) if row["parent_id"] else None,
row["role"],
content_json,
row["attachments_json"],
metadata_json,
int(row["created_at"]),
)
)
conn.executemany(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json,
metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
fork_messages,
)
for row in ancestry:
_replace_chat_attachment_inventory(
conn,
id_map[row["id"]],
row["attachments_json"],
row["content_json"],
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
thread_row = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,)
).fetchone()
return _chat_thread_from_row(thread_row) if thread_row is not None else None
except Exception:
conn.rollback()
raise
finally:
conn.close()
def count_forks_for_message(thread_id: str, message_id: str) -> int:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT COUNT(*) FROM chat_threads
WHERE forked_from_thread_id = ? AND forked_from_message_id = ?
""",
(thread_id, message_id),
).fetchone()
return int(row[0]) if row is not None else 0
finally:
conn.close()
def chat_thread_has_messages(thread_id: str) -> bool:
"""Whether this thread has any saved message. Existence only, no rows hydrated.
A temporary (incognito) chat is never written here, so this is what tells a thread
whose turns can be archived from one whose turns must not be.
"""
conn = get_connection()
try:
row = conn.execute(
"SELECT 1 FROM chat_messages WHERE thread_id = ? LIMIT 1", (thread_id,)
).fetchone()
return row is not None
finally:
conn.close()
def fork_counts_for_thread(thread_id: str) -> dict[str, int]:
"""Fork counts for every message of one thread, keyed by message id."""
conn = get_connection()
try:
rows = conn.execute(
"""
SELECT forked_from_message_id, COUNT(*) FROM chat_threads
WHERE forked_from_thread_id = ? AND forked_from_message_id IS NOT NULL
GROUP BY forked_from_message_id
""",
(thread_id,),
).fetchall()
return {str(row[0]): int(row[1]) for row in rows}
finally:
conn.close()
def list_chat_messages(thread_id: str) -> list[dict]:
conn = get_connection()
try:
rows = conn.execute(
"""
SELECT * FROM chat_messages
WHERE thread_id = ?
ORDER BY created_at ASC, id ASC
""",
(thread_id,),
).fetchall()
return [_chat_message_from_row(row) for row in rows]
finally:
conn.close()
def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT * FROM chat_messages
WHERE thread_id = ? AND id = ?
""",
(thread_id, message_id),
).fetchone()
return _chat_message_from_row(row) if row is not None else None
finally:
conn.close()
def _blob_part_base64_len(part: dict) -> int:
"""Base64 payload length of an image or audio content part, or 0."""
image = part.get("image")
if isinstance(image, str) and image[:5].lower() == "data:":
return len(image.rsplit(",", 1)[-1])
audio = part.get("audio")
if isinstance(audio, str) and _is_locally_stored_blob(audio):
return len(audio.rsplit(",", 1)[-1])
if isinstance(audio, dict):
data = audio.get("data")
if isinstance(data, str) and _is_locally_stored_blob(data):
return len(data)
return 0
def _attachment_content_parts(attachment: dict) -> list[dict]:
content = attachment.get("content")
if not isinstance(content, list):
return []
return [part for part in content if isinstance(part, dict)]
def _chat_attachment_size_bytes(attachment: dict) -> Optional[int]:
"""Approximate stored size of one attachment's content parts.
Image and audio parts hold base64 payloads (decoded bytes ~= 3/4 of the
encoded length); text parts count their character length. None when there
is no sizable content (e.g. a stripped/legacy attachment).
"""
total = 0
found = False
for part in _attachment_content_parts(attachment):
blob_len = _blob_part_base64_len(part)
if blob_len > 0:
total += (blob_len * 3) // 4
found = True
continue
text = part.get("text")
if isinstance(text, str) and text:
total += len(text.encode("utf-8", errors = "ignore"))
found = True
return total if found else None
def _content_part_attachments(content_json: Optional[str]) -> list[dict]:
"""Managed local blobs stored in content_json, with stable payload ids.
Exact duplicate blobs intentionally share one inventory id. Deleting that
id removes every identical copy, avoiding ambiguous index-based addressing.
"""
content = _json_loads(content_json, None)
if not isinstance(content, list):
return []
out: list[dict] = []
seen: set[str] = set()
for part in content:
if not isinstance(part, dict):
continue
attachment_id = _content_part_id(part)
payload = _managed_content_part_payload(part)
if attachment_id is None or payload is None or attachment_id in seen:
continue
seen.add(attachment_id)
kind, value = payload
content_type = None
part_name = part.get("name")
if isinstance(value, str) and value[:5].lower() == "data:":
content_type = value[5:].split(";", 1)[0].split(",", 1)[0] or None
out.append(
{
"id": attachment_id,
"type": kind,
"name": part_name
if isinstance(part_name, str) and part_name
else ("Chat image" if kind == "image" else "Chat audio"),
"contentType": content_type,
"content": [part],
}
)
return out
def list_chat_attachments_page(
limit: int = 50, offset: int = 0
) -> tuple[list[dict], Optional[int]]:
"""One bounded page from the normalized attachment inventory."""
if not 1 <= limit <= 100:
raise ValueError("limit must be between 1 and 100")
if offset < 0:
raise ValueError("offset must be non-negative")
conn = get_connection()
try:
_ensure_chat_attachment_inventory_current(conn)
rows = conn.execute(
"""
SELECT i.attachment_id, i.name, i.type, i.content_type,
i.size_bytes, m.id AS message_id, m.thread_id,
m.created_at, t.title AS thread_title, t.pair_id
FROM chat_attachment_inventory i
JOIN chat_messages m ON m.id = i.message_id
LEFT JOIN chat_threads t ON t.id = m.thread_id
ORDER BY m.created_at DESC, m.id ASC, i.attachment_id ASC
LIMIT ? OFFSET ?
""",
(limit + 1, offset),
).fetchall()
finally:
conn.close()
has_more = len(rows) > limit
page_rows = rows[:limit]
attachments = [
{
"id": row["attachment_id"],
"messageId": row["message_id"],
"threadId": row["thread_id"],
"pairId": row["pair_id"],
"threadTitle": row["thread_title"],
"name": row["name"],
"type": row["type"],
"contentType": row["content_type"],
"sizeBytes": row["size_bytes"],
"createdAt": row["created_at"],
}
for row in page_rows
]
return attachments, offset + limit if has_more else None
def list_chat_attachments() -> list[dict]:
"""Compatibility helper returning the full normalized inventory."""
attachments: list[dict] = []
offset = 0
while True:
page, next_offset = list_chat_attachments_page(limit = 100, offset = offset)
attachments.extend(page)
if next_offset is None:
return attachments
offset = next_offset
def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]:
"""One attachment record (full content) from a message, or None."""
conn = get_connection()
try:
row = conn.execute(
"""
SELECT message.attachments_json, message.content_json,
EXISTS(
SELECT 1 FROM chat_attachment_tombstones tombstone
WHERE tombstone.thread_id = message.thread_id
AND tombstone.message_id = message.id
AND tombstone.attachment_id = ?
) AS tombstoned
FROM chat_messages message
WHERE message.id = ?
""",
(attachment_id, message_id),
).fetchone()
finally:
conn.close()
if row is None or row["tombstoned"]:
return None
attachments = _json_loads(row["attachments_json"], None)
if isinstance(attachments, list):
for attachment in attachments:
if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id:
return attachment
if attachment_id.startswith(_CONTENT_PART_ID_PREFIX):
for attachment in _content_part_attachments(row["content_json"]):
if attachment["id"] == attachment_id:
return attachment
return None
def _record_chat_attachment_tombstone(
conn: sqlite3.Connection, thread_id: str, message_id: str, attachment_id: str
) -> None:
conn.execute(
"""
INSERT INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
VALUES (?, ?, ?, ?)
ON CONFLICT(thread_id, message_id, attachment_id) DO UPDATE SET
deleted_at = excluded.deleted_at
""",
(
thread_id,
message_id,
attachment_id,
int(datetime.now(timezone.utc).timestamp() * 1000),
),
)
def delete_chat_attachment(message_id: str, attachment_id: str) -> bool:
"""Remove one stored upload from a message.
The tombstone is retained while the thread exists, so pruning and later
recreating the same message id cannot restore the deleted upload. If an
ordinary attachment id collides with a content-blob id, both are deleted as
one managed item.
"""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
row = conn.execute(
"""
SELECT thread_id, attachments_json, content_json
FROM chat_messages WHERE id = ?
""",
(message_id,),
).fetchone()
if row is None:
conn.rollback()
return False
if str(message_id) in _research_message_ids(conn, str(row["thread_id"])):
conn.rollback()
raise ChatMessageProtectedError(
"Research prompts and responses are server-managed and cannot be edited"
)
attachments = _json_loads(row["attachments_json"], None)
updated_attachments_json = row["attachments_json"]
deleted_attachment = False
if isinstance(attachments, list):
remaining_attachments = [
attachment
for attachment in attachments
if not (
isinstance(attachment, dict)
and str(attachment.get("id") or "") == attachment_id
)
]
deleted_attachment = len(remaining_attachments) != len(attachments)
if deleted_attachment:
updated_attachments_json = json.dumps(remaining_attachments)
content = _json_loads(row["content_json"], None)
updated_content_json = row["content_json"]
deleted_content = False
if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list):
remaining_content = [
part
for part in content
if not (isinstance(part, dict) and _content_part_id(part) == attachment_id)
]
deleted_content = len(remaining_content) != len(content)
if deleted_content:
updated_content_json = json.dumps(remaining_content)
if not deleted_attachment and not deleted_content:
conn.rollback()
return False
conn.execute(
"""
UPDATE chat_messages
SET attachments_json = ?, content_json = ?
WHERE id = ?
""",
(updated_attachments_json, updated_content_json, message_id),
)
_record_chat_attachment_tombstone(
conn,
row["thread_id"],
message_id,
attachment_id,
)
_replace_chat_attachment_inventory(
conn,
message_id,
updated_attachments_json,
updated_content_json,
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return True
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]:
if not thread_ids:
return []
unique_thread_ids = list(dict.fromkeys(thread_ids))
messages: list[dict] = []
conn = get_connection()
try:
for start in range(0, len(unique_thread_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_thread_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT * FROM chat_messages
WHERE thread_id IN ({placeholders})
ORDER BY created_at ASC, id ASC
""",
chunk,
).fetchall()
messages.extend(_chat_message_from_row(row) for row in rows)
return sorted(
messages,
key = lambda message: (message["createdAt"], message["id"]),
)
finally:
conn.close()
def get_app_setting(key: str, fallback = None):
conn = get_connection()
try:
row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone()
if row is None:
return fallback
return _json_loads(row["value_json"], fallback)
finally:
conn.close()
def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
if not settings:
return {}
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO app_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in settings.items()],
)
conn.commit()
rows = conn.execute("SELECT key, value_json FROM app_settings ORDER BY key").fetchall()
return {row["key"]: _json_loads(row["value_json"], None) for row in rows}
finally:
conn.close()
def upsert_app_setting_map_entry(
key: str,
entry_key: str,
entry_value: dict[str, Any] | None,
*,
fill_absent_fields: bool = False,
) -> dict[str, Any]:
"""Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued
app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other
sub-entries cannot drop each other's updates.
``fill_absent_fields`` writes only what is missing: the entry is created when
it is not there, and otherwise gains the fields it does not already hold while
every stored value is left exactly as it is. Nothing is ever deleted. The read
and the write share this transaction, so a caller that read the map earlier
cannot replace a value written since. Used by the one-time localStorage
backfill, whose contract is that the server copy is the newer authority: an
upgraded install can hold an entry with only the fields an older release knew,
while this browser holds the rest, and entry-level skipping would strand them.
"""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone()
current = _json_loads(row["value_json"], {}) if row else {}
if not isinstance(current, dict):
current = {}
if fill_absent_fields:
if not entry_value:
conn.rollback()
return current
stored = current.get(entry_key)
if isinstance(stored, dict):
# Stored values win field by field, so this only adds.
merged = {**entry_value, **stored}
if merged == stored:
conn.rollback()
return current
current[entry_key] = merged
else:
current[entry_key] = entry_value
elif entry_value:
current[entry_key] = entry_value
else:
current.pop(entry_key, None)
now = datetime.now(timezone.utc).isoformat()
conn.execute(
"""
INSERT INTO app_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
(key, json.dumps(current), now),
)
conn.commit()
return current
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_chat_settings() -> dict[str, Any]:
conn = get_connection()
try:
rows = conn.execute("SELECT key, value_json FROM chat_settings ORDER BY key").fetchall()
settings: dict[str, Any] = {}
for row in rows:
settings[row["key"]] = _json_loads(row["value_json"], None)
return settings
finally:
conn.close()
def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]:
if not settings:
return list_chat_settings()
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO chat_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in settings.items()],
)
conn.commit()
return list_chat_settings()
finally:
conn.close()
# Discriminated unions, not partial patches: merging a `thread` pick into a stored
# `kb` one keeps `kbId`, which the payload's thread variant forbids.
_ATOMIC_SETTING_KEYS = frozenset({"ragSource"})
def _deep_merge_settings(current: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]:
merged = dict(current)
for key, value in updates.items():
current_value = merged.get(key)
if (
key not in _ATOMIC_SETTING_KEYS
and isinstance(current_value, dict)
and isinstance(value, dict)
):
merged[key] = _deep_merge_settings(current_value, value)
else:
merged[key] = value
return merged
def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]:
"""Atomic read-merge-write under BEGIN IMMEDIATE so concurrent writers
cannot drop each other's updates."""
if not updates:
return list_chat_settings()
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
current, corrupt = _load_chat_settings_for_merge(conn)
# An atomic key carries its whole value, so it repairs a quarantined row
# rather than patching a base that is no longer there.
unsafe_partial_keys = [
key
for key, value in updates.items()
if key in corrupt and isinstance(value, dict) and key not in _ATOMIC_SETTING_KEYS
]
if unsafe_partial_keys:
conn.commit()
keys = ", ".join(sorted(unsafe_partial_keys))
raise CorruptSettingsError(
f"Cannot apply partial settings patch to corrupt key(s): {keys}"
)
merged = _deep_merge_settings(current, updates)
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO chat_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in merged.items()],
)
conn.commit()
return merged
except CorruptSettingsError:
raise
except Exception:
conn.rollback()
raise
finally:
conn.close()
# --- Legacy Dexie import ledger (recovery rationale in _ensure_schema's schema comment) ---
def list_chat_legacy_imports() -> list[str]:
"""Return the legacy_thread_id of every thread already imported."""
conn = get_connection()
try:
rows = conn.execute("SELECT legacy_thread_id FROM chat_legacy_imports").fetchall()
return [row[0] for row in rows]
finally:
conn.close()
def upsert_chat_legacy_imports(legacy_thread_ids: list[str]) -> tuple[int, int]:
"""Mark each given legacy thread id as imported. Idempotent.
Returns (accepted, inserted): count of deduped non-empty input ids, and
count of rows actually new. RETURNING lets callers tell first-time imports
from idempotent re-runs without an extra SELECT.
"""
ids = list(dict.fromkeys(tid for tid in legacy_thread_ids if tid))
if not ids:
return 0, 0
ts = int(datetime.now(timezone.utc).timestamp() * 1000)
conn = get_connection()
try:
inserted = 0
for tid in ids:
row = conn.execute(
"""
INSERT INTO chat_legacy_imports (legacy_thread_id, imported_at)
VALUES (?, ?)
ON CONFLICT(legacy_thread_id) DO NOTHING
RETURNING legacy_thread_id
""",
(tid, ts),
).fetchone()
if row is not None:
inserted += 1
conn.commit()
return len(ids), inserted
finally:
conn.close()