* Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm The OpenAI-compatible endpoints serve whichever GGUF is loaded and ignore the request model field, so an OpenAI client that changes model never reloads. Add an opt-in setting that, when a /v1 request names a downloaded local GGUF different from the loaded one, loads it before serving by reusing the existing /load path (its dedup, tensor fallback, and threading apply). Unknown names still serve the loaded model, so drop-in compatibility is preserved and no remote download is triggered. Also add an optional idle auto-unload (TTL keep-warm): a pure-ASGI middleware tracks in-flight inference requests so a stream is never unloaded mid-response, and a lifespan loop unloads the model after the configured idle seconds. Both settings default off and live in the app_settings store, exposed via GET/PUT /api/settings/openai-auto-switch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: variant-aware auto-switch, /v1/responses coverage, keep-warm load stamp Follow-ups from review of the opt-in OpenAI auto-switch path: 1. Variant-aware dedup. _maybe_auto_switch_model compared only the repo id, so requesting another quant of the loaded repo (e.g. Q4_K_M loaded, Q8_0 asked) was served by the old quant. Compare hf_variant too, matching /load dedup. 2. Streaming /v1/responses now calls the auto-switch hook. It went straight into _responses_stream and only checked is_loaded, so stream=True could serve the old model or 400. Non-streaming already routed through chat completions; the hook is idempotent once loaded. 3. resolve_local_gguf tries an exact id match before splitting a trailing :VARIANT, so local ids that contain a colon (e.g. a Windows path) resolve instead of being cut at the drive letter. 4. Idle keep-warm stamps activity on a load/swap transition. _last_active was only refreshed by inference requests, so a model loaded after the server sat idle past the TTL could be unloaded before its first request. Tests cover each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make the /v1/responses auto-switch test order-independent The new streaming-responses test passed in isolation but failed under the CI's randomized collection order with "object has no attribute 'state'": it passed a bare object() as the request and stubbed only one dispatcher, so an ordering where the real dispatcher ran hit request.state. Give the request a state and stub both dispatchers; the test still asserts the hook fires before dispatch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: assert /v1/responses auto-switch wiring on source, not at runtime The behavioral version executed openai_responses and relied on stubbing its callees, which a randomized collection order in CI could defeat (the real dispatcher ran and hit request attributes). Assert on the function source that the hook precedes both dispatchers instead; the hook's runtime behavior is already covered by the direct _maybe_auto_switch_model tests. * Studio: auto-switch on /v1/embeddings, GGUF-only targets, idle-unload race gate Second-pass review follow-ups on the opt-in auto-switch path: 1. /v1/embeddings now calls the auto-switch hook before the loaded-state check, matching the other model-bearing OpenAI endpoints (the keep-warm middleware already treats embeddings as inference). 2. The resolver index is now GGUF-only. The local-model scanners also surface Transformers/safetensors repos; without a filter, auto-switch could unload the GGUF and route a request into the non-GGUF loader. _has_local_gguf checks a direct .gguf, a models-dir folder, and the HF-cache snapshots layout. 3. Idle keep-warm now holds an asyncio gate across the idle check and the unload, and a request bumps inflight under the same gate, so the loop can no longer unload in the window between "looks idle" and the kill. Tests cover each. Broader local-model source parity (LM Studio, Ollama, legacy caches, custom scan folders) is a follow-up; missing one of those today just falls through to the loaded model. * Studio: variant-aware local resolver, count_tokens + audio auto-switch coverage Third-pass review follow-ups on the opt-in auto-switch path: 1. The resolver is now variant-aware via list_local_gguf_variants. It indexes only the quants actually on disk, recursing snapshots and quant subdirs such as the nested per-quant folders, so a requested repo:VARIANT resolves only when that quant is local and a bare repo resolves to a concrete local quant. This fixes two gaps: the previous shallow glob rejected nested-variant GGUF repos, and a request for an uncached quant could send /load down the remote download path, breaking the local-only contract. 2. /v1/messages/count_tokens now auto-switches like its sibling /v1/messages, so a count uses the requested model's tokenizer. 3. /api/inference/audio/generate (direct GGUF TTS) is now tracked as in-flight inference, so the idle loop cannot unload the model mid-generation. Tests cover each. Two reviewer items are left as follow-ups: indexing the remaining local sources (LM Studio, Ollama, legacy/default caches, custom scan folders), which fails safe today by falling through to the loaded model; and fully serializing concurrent different-model requests, an inherent limit of the single-slot llama backend that the opt-in feature is not designed around. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make local GGUF resolver fail-safe so a bad model name cannot 500 The auto-switch hook calls resolve_local_gguf without its own guard, and /v1/completions and /v1/embeddings pass body.get("model") through unchanged. A non-string model (e.g. {"model": 123}) or any internal scan failure would then raise out of the resolver and turn a request that would otherwise be served by the loaded model into a 500, breaking the drop-in compatibility the feature is built on. Guard the resolver at its boundary: reject non-string input up front and wrap the lookup so any failure returns None (fall through to the loaded model). Add regression tests for both paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: per-model launch flags for auto-switched GGUF models * Studio: list switch-eligible GGUFs in /v1/models when auto-switch is on * Studio: settings UI for OpenAI model auto-switch and idle auto-unload * Studio: show save error over the disabled-idle hint in auto-switch settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address gemini review (case-insensitive /v1/models retrieve, idle-input empty guard) * Studio: address codex review (deterministic override args, exclude probe/embedding models from discovery) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep-warm count_tokens, gate idle on auto-switch, drop hidden models Three hardening fixes to the opt-in auto-switch path surfaced while reviewing the work that builds on it: 1. count_tokens keep-warm. /v1/messages/count_tokens counts via the loaded tokenizer and already auto-switches, but the keep-warm middleware did not track it, so idle auto-unload could free the model mid-count. It is now a tracked in-flight path. 2. "Off means unchanged" for idle unload. get_auto_unload_idle_seconds now reports 0 while auto-switch is disabled. Idle unload only makes sense with auto-switch on (an unloaded model returns only via the next request's swap), so a stray TTL can no longer trigger a destructive unload while the feature is off, keeping the disabled state identical to pre-feature behavior. 3. Hidden models are not switch targets. The resolver index now skips what Studio hides from its own pickers (the llama.cpp validation probe, RAG embedding weights) via _is_hidden_model, so they can never be auto-switched to by name. Tests added for each. * Studio: bare-id reuse, responses validation order, in-flight tracking Review follow-ups after folding in the per-model overrides and discovery work: 1. A bare model id (no :VARIANT) is now satisfied by any loaded quant of that repo. Previously a bare name resolved to the largest local quant, so it could force a slow reload when a different quant of the same repo was already serving. An explicit repo:VARIANT request still honors the quant. 2. /v1/responses now runs the auto-switch hook after the empty-input validation so a request that 400s can no longer trigger a multi-minute model load before being rejected. The hook still precedes both dispatchers, so streaming requests switch. 3. The keep-warm middleware now tracks in-flight requests whenever auto-switch is enabled rather than only when the idle TTL is already positive, so a stream that starts with the TTL at 0 is still protected if idle-unload is enabled mid-stream. Off still passes straight through. Tests added for each. * Studio: tighten auto-switch code comments Comment/docstring-only pass over the OpenAI auto-switch feature: collapse multi-line blocks, drop a comment that restated the gate it sits next to, and trim verbose docstrings on internal helpers while keeping the load-bearing rationale (concurrency, API behavior, drop-in compat, gotchas). No logic change: verified comment-only with the AST/printer signature check. * Studio: bind auto-switch locks per running loop Review follow-up. The auto-switch swap lock and the keep-warm unload gate were module-level asyncio.Lock objects. That is safe under the single uvicorn loop and on Python 3.10+ (the Lock resolves the running loop lazily on acquire), but a module-level Lock binds to one loop on pre-3.10, which can raise a loop mismatch in multi-loop runners. Resolve each lock through a per-loop accessor backed by a WeakKeyDictionary so every running loop gets its own Lock and stale loops are collected. No behavior change under the server's single loop. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch re-review fixes (body codes, coverage, swap, alias, tracking) Follow-ups from a second review pass over the opt-in OpenAI auto-switch feature: 1. OFF-state status codes: /v1/completions and /v1/embeddings moved the body read ahead of the loaded-state check, so a malformed/empty body with no model loaded returned 500 instead of the prior 503. A shared helper reads the body defensively (an unparseable/non-dict body yields no model), and the handler re-reads after the 503 gate to surface the original parse error exactly as before. OFF behavior is unchanged. 2. Local-model coverage: the resolver index only scanned ./models and the active HF cache, while the model picker also lists the legacy/default HF caches, LM Studio dirs, and user scan folders. A request for one of those named models silently served the loaded model instead. _build_index now scans the same roots (Ollama's symlink-creating scanner is skipped on the request path), and resolution is offloaded with asyncio.to_thread so the wider scan never blocks the event loop. 3. Swap vs in-flight stream: a cross-model swap killed the llama-server while another client was still streaming from it. The hook now tracks how many requests are streaming on the loaded model (in-flight minus those still inside the hook) and returns 409 instead of swapping while one is active. Concurrent same-model requests never reach this path, so they are unaffected. 4. Idle-unload + alias: after idle-unload freed the model, an unknown/alias name resolved to nothing and 503'd, though it served the active model before the TTL. Idle-unload now remembers the freed id and an alias request reloads it (only an already-local model, so no remote download), cleared once a model is loaded again. 5. In-flight tracking: the keep-warm middleware tracked in-flight only while the feature was on, so a stream started while off could be unloaded if idle-unload was enabled mid-stream. It now tracks on every inference path; counting is cheap and invisible to clients. Tests added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove stray async_task_outputs files committed by mistake * Studio: auto-switch review round 3 (revert swap guard, hardening) Addressing a third review pass: - Revert the cross-model swap guard. It counted keep-warm in-flight (which includes external-provider calls that never touch the local model) and so could 409 a local swap spuriously, and it still left a same-model request able to start streaming on the model a concurrent swap was unloading. A correct fix needs a request-lifetime reader/writer barrier; a partial guard was worse than the honest single-slot behavior, so concurrent different-model use is back to being serialized (documented), like llama-swap's single slot. - Non-string request model (e.g. {"model": 123} on a raw-body endpoint) is now treated as absent, so it falls through instead of raising in the membership checks once an idle-unload stash exists. - Idle-unload now stashes and replays the freed quant: an alias reload restores the exact (id, variant) that was freed rather than the largest local quant. - Anthropic /v1/messages validates max_tokens before the auto-switch hook, so a request that 400s never triggers a model load. - Keep-warm tracks a pending count for requests waiting on the unload gate, so the idle loop cannot unload the model out from under a request that is blocked on the gate but not yet counted as in-flight. - The idle-unload task is awaited after cancel on shutdown to avoid pending-task warnings. - The resolver's HF cache scan is None-safe and logs at debug instead of letting a bad root abort the whole index build. - upsert_app_setting_map_entry rolls back explicitly on error. Tests updated/added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep saved idle-unload seconds when auto-switch is toggled off * Studio: auto-switch hardening (thread-safe lock maps, body validation) Defensive fixes from review: - Guard the per-loop WeakKeyDictionary get-or-create for both the unload gate and the auto-switch lock with a threading lock, since WeakKeyDictionary mutation is not thread-safe when two event loops run on different threads. - Build the resolver index under the cache lock so concurrent callers with an expired cache don't all run the multi-dir scan at once. - /v1/completions and /v1/embeddings return a clean 400 for a valid JSON body that is not an object (e.g. a list), instead of a 500 from body.get(...). - The keep-warm middleware only tracks POST requests (inference is always POST), so CORS preflight (OPTIONS) is not counted, and tolerates a None path. Tests added for the list-body 400 and the non-POST skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 4 (local-path load, swap guard, idle fixes) From a 10-reviewer pass: - HF-cache entries now load by a concrete local path, not the bare repo id. The resolver records a load_path (the snapshot dir for a models--* cache repo, the file/dir otherwise) so /load takes the local branch and can never trigger a download to satisfy a partial cache. The advertised loader_id (repo id) is kept as the launch-override key. resolve_local_gguf now returns (load_path, variant, loader_id). - Re-add a single-slot swap guard: a cross-model swap returns 409 model_switch_busy while another inference request is active rather than killing its stream (the caller is excluded from the count), and holds the keep-warm gate across the load so no new inference starts mid-swap. Concurrent same-model requests never reach this path. A residual spurious 409 is possible while a concurrent or external- provider request is active; that is the documented single-slot tradeoff. - Idle keep-warm tracks (model_identifier, hf_variant): reloading the same repo at a different quant counts as a fresh model, so it is not unloaded before one TTL. - Track Studio's own /api/inference/generate/stream so the idle loop can't unload the model mid-stream on that route. - A successful manual /load clears the idle-unload reload stash synchronously, not only on the next idle poll. Also merged origin/main (the branch had fallen behind, which would have reverted unrelated files on merge). Tests added/updated for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 5 (concurrency, identity, load gate) From a 10-reviewer pass (9 request-changes, 1 approve): - Concurrent same-target requests load once instead of each returning 409. The count-based busy guard could not tell "another request wants the same model" (safe, load once) from "another request is using the loaded model" (refuse). Track in-flight auto-switch requests per (target, variant) and subtract same-target waiters from the busy count; a cross-model swap still 409s while a genuinely different request is active. - Fix the identity confusion introduced when round 4 began loading by concrete local path: the backend identifier became a filesystem path. Record the advertised repo id on the backend after an auto-switch load and use it so (a) a model loaded manually by repo id is recognized as already serving (no spurious reswap/409), (b) /v1/models reports the repo id, never a host path or a duplicate, and (c) the idle-unload stash keeps the override keyed by the repo id, so an alias reload after TTL keeps the user's saved launch flags. - Gate the manual /load route with the keep-warm lifecycle gate so idle auto-unload can't unload a model mid-load. load_model now wraps _load_model_impl in the gate; auto-switch calls _load_model_impl directly since it already holds the gate. - Restore default-off parity on Anthropic /v1/messages: an unloaded backend with auto-switch disabled 503s before the max_tokens 400 check, as it did pre-feature. When the feature is on, request-shape validation still runs before any load. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 6 (concurrency ordering, leaks, unload gate) From a second 10-reviewer pass (8 request-changes, 2 approve): - Same-target concurrency: register a waiter by the raw requested model before the (slow) resolve, and exclude pending requests from the swap busy count. The middleware counts a concurrent same-model request as in-flight before it resolves and joins the resolved-target waiter map, so the prior fix could still 409 it. The guard now subtracts max(same resolved-target, same raw-request) waiters and ignores pending (a pending request is blocked in the middleware, not generating, so a swap can't interrupt it). - External-provider requests no longer block a local swap. The keep-warm middleware counts every inference-path POST, but external-provider chat returns before the auto-switch hook and never touches the local GGUF. The chat handler now untracks itself before proxying, so its in-flight stream can't trip model_switch_busy on a concurrent local auto-switch. The middleware skips its own end-decrement for an untracked request. - Manual /unload is gated like load and idle-unload: it holds the lifecycle gate and returns 409 rather than tearing down llama-server while an inference request is in flight. - Response model id no longer leaks the load path. /v1/models already advertised the repo id; chat, completions, embeddings, Anthropic messages, and audio response bodies now use the same _llama_public_model_id helper instead of the concrete on-disk model_identifier. - Chat completions validates the non-system-message requirement before the auto-switch hook (as /responses and /messages already do), so an invalid request can't swap the resident model before returning 400. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 7 (teardown policy, Unsloth-active swap, training) From a third 10-reviewer pass (9 request-changes, 1 approve), all on the same asymmetric-teardown theme. Resolved per the intended policy that only automatic paths defer to an active stream; deliberate user actions stay interrupting: - Revert the manual /unload in-flight guard added last round. A manual /load or /unload is a deliberate action and tears down immediately, as before; only the automatic idle-unload loop and auto-switch defer to an active request. This removes the asymmetry the reviewers flagged (manual /load, the /unload Unsloth branch, and the opposite-backend swaps inside _load_model_impl) by not extending the guard to deliberate paths, rather than spreading it. - Auto-switch now refuses a swap whenever another inference request is in flight, not only when a GGUF is already loaded. _load_model_impl also unloads an active Unsloth/transformers backend before loading a GGUF, so the busy guard must cover that case too; otherwise an Unsloth stream could be killed by an auto-switch. - Refuse API-initiated training while inference is active. When Studio is driven as an inference API (sk-unsloth key auth), POST /api/training/start returns 409 if a request is in flight, since training frees VRAM by unloading the chat model and would kill the stream. The Studio UI (session auth) still starts training and coexists/frees VRAM as before. A mixed UI+API session is not yet special-cased. Adds auth.authentication.authenticated_via_api_key. Tests added/updated for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add UNSLOTH_MODEL_IDLE_TTL env override for idle-unload Borrowed from PR 6517: a startup env var that sets the idle-unload TTL without the settings UI. Unlike the stored setting (gated on auto-switch), the env value is a standalone default that enables idle-unload even with auto-switch off, for headless/container deploys. An explicit UI/API value still overrides it and stays gated. The settings GET reflects the env default when nothing is stored. * Studio: auto-switch fixes from review (paths, embeddings input, env idle reload) - /v1/models advertises a client-facing alias instead of a filesystem path: the ./models and LM Studio scanners report the on-disk path as the model id, so the index now prefers model_id/display_name as the advertised/override id and keeps the concrete path internal as load_path, still resolvable by path. - /v1/embeddings validates input before auto-switch: a request with a model but no input now 400s before the hook (like chat/responses/messages), so an invalid embeddings request cannot unload or swap the resident model. - Standalone UNSLOTH_MODEL_IDLE_TTL reloads the freed model: the hook now runs when auto-switch or idle-unload is active, and with auto-switch off it skips the resolver and only restores the idle-unloaded model, so the first idle timeout no longer leaves later /v1 requests with nothing loaded. - Do not resurrect a stale GGUF over an active Unsloth model: the reload-stash path bails when a non-GGUF backend is loaded, so an unknown /v1 name cannot tear down a live Transformers/Unsloth model. - Defensive HF cache scan: each cache root's resolve/dedup is wrapped so a missing or malformed root skips that root rather than aborting the index. - Single-model retrieve checks the id is a string before lowercasing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix automatic-load asymmetry, audio reload, preview, idle timer The standalone UNSLOTH_MODEL_IDLE_TTL reload is a second automatic-load trigger, but several validate-before-switch guards and reload hooks only checked the auto-switch toggle. Add a shared _automatic_model_load_may_run() (auto-switch on, or idle TTL > 0) and route every guard through it. - /v1/completions validates prompt before any automatic load (it was the one model-bearing route with no pre-check). - /v1/chat/completions and /v1/embeddings pre-checks gate on the shared predicate so a standalone idle TTL cannot reload then reject. - /v1/messages no longer 503s before the reload hook can restore an idle-freed model when auto-switch is off. - Raw completions/embeddings with no model field pass a non-empty sentinel so the idle-stash reload runs, restoring the legacy "omit model, use loaded" path. - /api/inference/audio/generate gains the reload hook (after message validation) so an idle-freed audio GGUF is restored. - Public preview opts out of auto-switch via a request-scope flag, so a caller's model field cannot swap away from the pinned checkpoint; preview chat streams are now matched by _is_inference_path so idle-unload cannot kill them. - Keep-warm no longer stamps activity on request start, and external-provider untracking decrements without restamping, so periodic external traffic can no longer keep the local GGUF warm forever. Merges origin/main (the branch had fallen behind, which also brought in the preview route the review flagged). * Studio: surface model auto-switch in the API tab and demo it in examples The OpenAI auto-switch toggle previously lived only in Settings -> General. Add the same toggle to the API tab's usage-examples panel (it shares the settings cache), and make the examples reflect it: when on, the Python examples append a second call naming a different downloaded GGUF (so the model field visibly selects which model serves), and the curl examples gain a one-line note. Reuses the existing settings API client and i18n keys. * Studio: harden OpenAI auto-switch reload-only path and Anthropic tool validation - Omitted-model raw-body requests pass a reload-only sentinel so the idle-stash reload still restores an idle-freed model, but the resolver never matches a downloaded GGUF literally named "default". - Reject malformed Anthropic client tools before _maybe_auto_switch_model so an invalid request can no longer evict the loaded model. * Studio: extend auto-switch reload-only and tool validation to schema endpoints - Schema-backed endpoints (chat completions, responses, count_tokens, messages, audio) defaulted an omitted model to "default" and passed it to the switch hook, so a downloaded GGUF named "default" could be swapped to. Route the hook through a helper that switches only on an explicitly set model, else reload-only. - Propagate the explicit-set status when building the chat request from a Responses request, so the non-streaming chat re-check stays reload-only too. - Validate Responses function tools before the switch hook so a malformed tool returns 400 without evicting the loaded model. * Studio: serialize auto-switch swaps across event loops with a process-wide gate The auto-switch lock is a per-event-loop asyncio.Lock, so two /v1 swaps on different loops in one process could both pass it and race the single model slot (the backend and _load_model_impl are process-wide). Add a process-wide threading gate around the swap, acquired off the loop so a cross-loop wait never blocks it, layered with the existing per-loop lock. Add a cross-loop test that fails without the gate (two slow loads overlap) and passes with it. * Studio: make the auto-switch swap gate wait cancellation-safe _acquire_swap_gate awaited asyncio.to_thread(lock.acquire) when another loop held the process-wide gate. to_thread cancellation doesn't stop the worker thread, so a /v1 request cancelled mid-wait (client disconnect during a cross-loop swap) would have its thread acquire the gate after the fact, while the finally that releases it never runs -- permanently deadlocking later auto-switch swaps. Poll a non-blocking acquire off a short asyncio.sleep instead: it still keeps the wait off the loop and serializes across loops, but a cancel now lands during the sleep, when the gate is not held, so nothing leaks. Add a test that deadlocks the to_thread variant (it times out) and passes with the poll. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: validate modality and tool-confirmation before auto-switch Two more request shapes could load a named GGUF and only then 400, evicting the resident model: - An image request naming a different text-only GGUF. The switch hook now takes require_vision and rejects a swap to a non-vision target before loading it; a GGUF's vision capability is its companion mmproj, knowable without a load, and matches the post-load guard. Only the resolver branch is checked, never the reload-stash restore. - confirm_tool_calls=true with stream=false and local tools. /v1/chat/completions now rejects that shape before the hook, mirroring the local tool path's bypass_permissions exemption and intent signal. The vision probe threads the ambient HF token to keep the capability-probe invariant. Reload-only and idle-reload paths are unaffected. * Studio: extend validate-before-switch and make the lifecycle gate process-wide - /v1/messages/count_tokens now rejects malformed client tools before the switch hook, like /messages (shared _validate_anthropic_client_tools helper), so a count request can't evict the loaded model. - /v1/chat/completions rejects a malformed tool_choice forcing object (a {"type":"function","function":{}} with no name) before the switch hook. - The inference lifecycle gate that blocks new inference during a swap is now process-wide (a poll-acquired threading lock, cancellation-safe), not a per-loop asyncio lock, so a request on another event loop can't start inference while a swap tears the single backend down. - Usage examples no longer hard-code a switch-demo repo most users lack; the model is an explicit placeholder the user replaces. * Studio: extend the auto-switch modality guard to /v1/responses and /v1/messages The pre-load vision check that guards /v1/chat/completions now also runs on /v1/responses and /v1/messages, so an image request naming a text-only GGUF is rejected before the swap and never evicts the resident vision model. Run the vision capability probe off the event loop. Make the /v1/models retrieve loaded fast-path case-insensitive, and never advertise a host path from the resolver. Remove the dead list_switch_eligible_ids helper, superseded by the /v1/models catalog. * Studio: filter /v1/models to GGUF, per-loop catalog lock, reject system-only Responses Address review findings on the auto-switch path: - /v1/models advertises only GGUF models the API can actually switch to; a safetensors/LoRA entry would be selectable but never loadable via llama.cpp. - The /v1/models catalog cache uses a per-loop lock (like the auto-switch path) so a second event loop awaiting it can't hang in a multi-loop process. - /v1/responses rejects system/developer-only input before the switch, mirroring chat, so an invalid request can't evict the resident model. - _build_index guards each scan source on its own so one bad root drops only that source; the vision probe logs a real detection failure instead of swallowing it. * Studio: list cached GGUFs in /v1/models by inspecting files, not model_format The HF-cache scanner leaves model_format unset for GGUF snapshots, so the previous model_format == "gguf" filter dropped every downloaded HF-cache GGUF from /v1/models and the retrieve fallback. Decide GGUF-ness from the on-disk files via the resolver (info_has_local_gguf) instead, run off the event loop, so the catalog advertises exactly what /v1 can serve. * Studio: fix /v1/messages/count_tokens route binding plus auto-switch review fixes The @router.post decorator for /messages/count_tokens had been separated from anthropic_count_tokens by the _validate_anthropic_client_tools helper, so the route bound to the validator and dropped its auth dependency. Move the decorator back onto the handler. Add route-binding tests asserting each /v1 endpoint maps to its handler with the auth dependency, so a decorator/handler split is caught at the route level (the direct-call tests missed it). Also from review: - update_openai_auto_switch writes both settings keys in one transaction so a PUT can't leave one updated and the other stale (drop the now-unused single setters). - max_seq_length override rejects 0 at the boundary (ge=1) instead of accepting then silently dropping it. - Document that embeddings auto-switch is best-effort: GGUF pooling has no cheap pre-load probe like vision's mmproj, so a guard would false-reject GGUF embedders. - Add a positive idle-unload test (loop frees the model and stashes it for reload). * Studio: validate Responses tool_choice + Anthropic mixed tools before switch, filter Ollama from catalog More auto-switch review findings: - /v1/responses rejects a forcing-function tool_choice with no name before the switch, mirroring chat, so a malformed request can't evict the resident model. - /v1/messages rejects mixing Anthropic server tools with custom client tools before the switch (the check depends only on the payload, so it moves up cleanly). - /v1/models no longer advertises Ollama-link models: info_has_local_gguf excludes .studio_links / ollama_links entries, which the resolver skips and can't switch to, so an advertised id never silently falls through. * Studio: guard chat audio input before switch; surface env-backed idle unload in settings UI A chat request carrying audio_base64 rides the same companion mmproj projector as a vision request, so a text-only target cannot serve it either. Flag require_vision for audio input as well so the multimodal probe runs before the switch and a rejected request never evicts the working model. Generalize the reject message to cover image and audio. The settings response now reports idle_unload_active (effective TTL > 0) so the UI can distinguish idle-unload that is active via the UNSLOTH_MODEL_IDLE_TTL env var from the case where it needs the toggle enabled. * Studio: harden auto-switch eviction guards (count_tokens vision, TTS reload-only, mmproj/stash) Four eviction/correctness fixes on the opt-in /v1 auto-switch path: - /v1/messages/count_tokens now carries the same require_vision guard as /messages, so an image count naming a text-only GGUF can't evict a loaded vision model for a swap that can't serve the request. - /audio/generate is now reload-only. A local GGUF's audio-input capability is not a cheap pre-load probe (the companion mmproj signal can't tell an audio projector from a vision one, and codec TTS ships no projector), so resolving the client model could load a text/vision-only target and evict the working audio model before the audio check fails. Only the idle-stash restore runs here; switching TTS models is an explicit /load. - The resolver no longer treats a standalone mmproj .gguf as a servable model. _scan_models_dir's standalone-file pass does not filter mmproj the way its directory scan does, so /v1/models could advertise a projector and a switch could load it over the real weights. - A non-GGUF (Transformers/Unsloth) load and a deliberate /unload now clear the idle reload stash, so a manual load/unload is never superseded by a stale idle-freed GGUF that the next /v1 request resurrects. * Studio: report advertised repo id consistently after an auto-switch Two model-id reporting fixes so an auto-switched cached HF GGUF is named by its repo id everywhere, not its snapshot path: - Streamed /v1/responses envelopes now derive the model id from _llama_public_model_id (which prefers _openai_advertised_id) instead of the raw model_identifier. After an auto-switch the identifier is the snapshot path while the repo id lives in _openai_advertised_id, so the stream used to report a snapshot basename while /v1/models, chat completions, and non-streaming Responses all reported the repo id. - When an advertised alias already resolves to the loaded model (a model loaded by local path, requested by its repo or LM Studio id), the already-serving early return now records the alias as the advertised id, so /v1/models and responses report the alias and mark it loaded instead of the path-derived basename. Resolver branch only; safe lock-free because an in-flight request blocks any concurrent swap via the single-slot busy guard. * Studio: validate request shapes before auto-switch (prompt/input/audio/mcp confirm) Four more validate-before-switch guards so a deterministic client error never evicts the resident model on the opt-in /v1 auto-switch path: - /v1/completions rejects an object/number prompt (only a string or array is valid) before the switch, instead of loading the named GGUF and letting llama-server reject the shape afterward. - /v1/embeddings rejects an object/number input the same way. - Chat rejects an oversized audio_base64 upload (413) before the switch. The size cap is a cheap, target-independent length check; the decode itself stays post-switch to avoid decoding a valid upload twice. - The chat confirm-without-stream pre-switch guard now mirrors the tool loop's actual enablement: _effective_enable_tools (honoring a CLI --enable-tools policy) and mcp_enabled (which opens the tool loop on its own but defers to a CLI --disable-tools policy). Previously a confirm+no-stream request with only mcp_enabled slipped past and 400'd after the swap. * Studio: fix model-id retrieval, streaming n>1, resolver cache TTL, keep-warm auth Four fixes from review: - GET /v1/models/{id} legacy raw-path fallback now maps the raw identifier to the same public id its /v1/models entry uses. After an auto-switch load the identifier is the snapshot path while the entry is keyed by the advertised repo id, so a client that cached the old absolute path no longer 404s on a model that is in fact loaded. - stream=true with n>1 is now rejected before the switch. Only the non-streaming GGUF path returns multiple choices, so streaming n>1 is invalid on every local serving path; both fields are known pre-switch, so it must not load model B only to 400 and evict model A. Non-streaming n>1 stays post-switch where the serving path decides. - The resolver index cache is stamped after _build_index, not with the pre-scan timestamp. On installs with enough local models for the multi-root scan to exceed the 5s TTL, the cache was stored already expired and every request rebuilt it. - The keep-warm middleware no longer stamps model activity for 401/403 responses. It runs before FastAPI auth, so unauthenticated probes used to refresh the idle timer without touching llama.cpp; they now decrement the in-flight count without keeping the model warm. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
1860 lines
61 KiB
Python
1860 lines
61 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 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 project_workspaces_root, studio_db_path, ensure_dir
|
|
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 []
|
|
|
|
|
|
_schema_lock = threading.Lock()
|
|
_schema_ready = False
|
|
_SQLITE_IN_CHUNK_SIZE = 900
|
|
_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",)
|
|
|
|
|
|
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)
|
|
|
|
|
|
def _ensure_project_workspace(root_path: str) -> str:
|
|
root = Path(root_path).expanduser()
|
|
root_resolved = ensure_dir(root).resolve()
|
|
for subdir in _PROJECT_WORKSPACE_SUBDIRS:
|
|
ensure_dir(root_resolved / subdir)
|
|
return str(root_resolved)
|
|
|
|
|
|
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 _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
|
|
)
|
|
"""
|
|
)
|
|
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")
|
|
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 keep
|
|
# case-sensitive BINARY so /Models and /models stay 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,
|
|
openai_code_exec_container_id TEXT,
|
|
anthropic_code_exec_container_id TEXT,
|
|
forked_from_thread_id TEXT,
|
|
forked_from_message_id 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 "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")
|
|
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
|
|
)
|
|
"""
|
|
)
|
|
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; Dexie is read-only so per-thread suffices.
|
|
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)"
|
|
)
|
|
|
|
|
|
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()
|
|
|
|
|
|
def get_connection() -> 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))
|
|
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)
|
|
_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],
|
|
) -> None:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(id, model_name, dataset_name, config_json, started_at, total_steps),
|
|
)
|
|
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,
|
|
) -> None:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute(
|
|
"""
|
|
UPDATE training_runs
|
|
SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
|
|
duration_seconds = ?, loss_sparkline = ?, output_dir = ?,
|
|
error_message = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
status,
|
|
ended_at,
|
|
final_step,
|
|
final_loss,
|
|
duration_seconds,
|
|
loss_sparkline,
|
|
output_dir,
|
|
error_message,
|
|
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 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,
|
|
CASE
|
|
WHEN r.status = 'stopped'
|
|
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')
|
|
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 = 'stopped'
|
|
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')
|
|
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 = 'stopped'
|
|
AND NOT EXISTS (
|
|
SELECT 1
|
|
FROM training_runs newer
|
|
WHERE newer.output_dir = r.output_dir
|
|
AND newer.status IN ('stopped', 'completed')
|
|
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("DELETE FROM training_runs WHERE id = ?", (id,))
|
|
conn.commit()
|
|
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 = 'error',
|
|
error_message = 'Server restarted during training',
|
|
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(path: str) -> dict:
|
|
"""Add a directory to the custom scan folder list. Returns the row."""
|
|
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")
|
|
if not os.access(normalized, os.R_OK | os.X_OK):
|
|
raise ValueError("Path is not readable")
|
|
|
|
# Windows: normcase for the denylist check but store original casing
|
|
# so consumers see the native drive-letter 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):
|
|
raise ValueError(f"Path under {prefix} is not allowed")
|
|
|
|
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)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO scan_folders (path, created_at) VALUES (?, ?)",
|
|
(normalized, now),
|
|
)
|
|
conn.commit()
|
|
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)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def remove_scan_folder(id: int) -> None:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,))
|
|
conn.commit()
|
|
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) -> dict:
|
|
data = dict(row)
|
|
return {
|
|
"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"],
|
|
"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"),
|
|
}
|
|
|
|
|
|
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
|
|
|
|
|
|
def upsert_chat_thread(thread: dict) -> dict:
|
|
conn = get_connection()
|
|
try:
|
|
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)
|
|
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,
|
|
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
|
|
""",
|
|
(
|
|
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"]),
|
|
thread.get("openaiCodeExecContainerId"),
|
|
thread.get("anthropicCodeExecContainerId"),
|
|
thread.get("forkedFromThreadId"),
|
|
thread.get("forkedFromMessageId"),
|
|
),
|
|
)
|
|
conn.commit()
|
|
return get_chat_thread(thread["id"]) or thread
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
|
|
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")),
|
|
"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"),
|
|
),
|
|
}
|
|
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_thread(id)
|
|
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute(
|
|
f"UPDATE chat_threads SET {', '.join(assignments)} WHERE id = ?",
|
|
(*values, id),
|
|
)
|
|
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 created_at DESC",
|
|
values,
|
|
).fetchall()
|
|
return [_chat_thread_from_row(row) for row in rows]
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_chat_threads(ids: list[str]) -> None:
|
|
if not ids:
|
|
return
|
|
conn = get_connection()
|
|
try:
|
|
conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids])
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def clear_chat_history() -> None:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute("DELETE FROM chat_threads")
|
|
conn.commit()
|
|
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)
|
|
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()
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
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)
|
|
conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,))
|
|
conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,))
|
|
conn.commit()
|
|
if delete_files:
|
|
_delete_project_workspace(project)
|
|
return project
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class ChatMessageConflictError(RuntimeError):
|
|
"""Raised when a chat message id already belongs to another thread."""
|
|
|
|
|
|
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 upsert_chat_message(message: dict) -> dict:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
_raise_if_chat_message_thread_conflicts(
|
|
conn,
|
|
message["threadId"],
|
|
[message["id"]],
|
|
)
|
|
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
|
|
""",
|
|
(
|
|
message["id"],
|
|
message["threadId"],
|
|
message.get("parentId"),
|
|
message["role"],
|
|
json.dumps(message.get("content", [])),
|
|
json.dumps(message.get("attachments"))
|
|
if message.get("attachments") is not None
|
|
else None,
|
|
json.dumps(message.get("metadata"))
|
|
if message.get("metadata") is not None
|
|
else None,
|
|
int(message["createdAt"]),
|
|
),
|
|
)
|
|
conn.commit()
|
|
return message
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def sync_chat_messages(
|
|
thread_id: str,
|
|
messages: list[dict],
|
|
prune_missing: bool = False,
|
|
) -> list[dict]:
|
|
conn = get_connection()
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
_raise_if_chat_message_thread_conflicts(
|
|
conn,
|
|
thread_id,
|
|
[m["id"] for m in messages],
|
|
)
|
|
if prune_missing:
|
|
conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,))
|
|
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"],
|
|
json.dumps(m.get("content", [])),
|
|
json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
|
|
json.dumps(m.get("metadata")) if m.get("metadata") is not None else None,
|
|
int(m["createdAt"]),
|
|
)
|
|
for m in messages
|
|
],
|
|
)
|
|
conn.commit()
|
|
return list_chat_messages(thread_id)
|
|
except ChatMessageConflictError:
|
|
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 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")
|
|
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)
|
|
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,
|
|
),
|
|
)
|
|
conn.executemany(
|
|
"""
|
|
INSERT INTO chat_messages
|
|
(id, thread_id, parent_id, role, content_json, attachments_json,
|
|
metadata_json, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
[
|
|
(
|
|
id_map[row["id"]],
|
|
new_thread_id,
|
|
id_map.get(row["parent_id"]) if row["parent_id"] else None,
|
|
row["role"],
|
|
row["content_json"],
|
|
row["attachments_json"],
|
|
row["metadata_json"],
|
|
int(row["created_at"]),
|
|
)
|
|
for row in ancestry
|
|
],
|
|
)
|
|
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 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 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
|
|
) -> 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."""
|
|
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 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()
|
|
|
|
|
|
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 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)
|
|
unsafe_partial_keys = [
|
|
key for key, value in updates.items() if key in corrupt and isinstance(value, dict)
|
|
]
|
|
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
|
|
# ---------------------------------------------------------------------------
|
|
# See the schema comment in _ensure_schema() for the recovery rationale.
|
|
|
|
|
|
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()
|