Merge main into feat/studio-parallel-slots-ui for PR #7447
Second pass, for the commits that landed after the previous merge (#7514, #7520, #7530, #7543). One conflict, in the validate preflight's load-kwargs call: main added cache_type_kv and tensor_parallel to it while this branch changed n_parallel to read the single resolved value. Keep both, so the preflight still sizes from _n_parallel and carries main's new arguments. test_active_generations.py and test_parallel_slots_per_load.py are 114 passed. tests/studio is 1861 passed with the 8 pre-existing MLX hardware-dispatch failures that main shows on this machine too.
This commit is contained in:
Daniel Han
commit
936ea4d072
38 changed files with 2839 additions and 416 deletions
3
.github/workflows/studio-frontend-ci.yml
vendored
3
.github/workflows/studio-frontend-ci.yml
vendored
|
|
@ -133,6 +133,9 @@ jobs:
|
|||
- name: Typecheck
|
||||
run: npm run typecheck
|
||||
|
||||
- name: Unit tests
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
run: npm run build
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -90,9 +90,10 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
|
|||
def _flag_name(token: str) -> Optional[str]:
|
||||
"""Flag name for ``token``, or None if it isn't a flag.
|
||||
|
||||
Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts
|
||||
always start with a letter), and normalises attached `-np8` / `-np-1` /
|
||||
`-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`.
|
||||
Peels `--key=value` to `--key`, normalises long-option underscores like
|
||||
llama.cpp, treats `-1`/`-0.5` as values (shorts always start with a letter),
|
||||
and normalises attached `-np8` / `-np-1` / `-np8x` to `-np`. Mirrors the
|
||||
CLI's `_expand_attached_np_short`.
|
||||
"""
|
||||
token = token.strip()
|
||||
if not token.startswith("-") or token in {"-", "--"}:
|
||||
|
|
@ -100,6 +101,8 @@ def _flag_name(token: str) -> Optional[str]:
|
|||
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
|
||||
return None
|
||||
name = token.split("=", 1)[0]
|
||||
if name.startswith("--"):
|
||||
name = name.replace("_", "-")
|
||||
if len(name) > 3 and name.startswith("-np"):
|
||||
suffix = name[3:]
|
||||
if suffix[0].isdigit() or (
|
||||
|
|
|
|||
|
|
@ -267,6 +267,8 @@ class ValidateModelRequest(BaseModel):
|
|||
# /load; defaults preserve old behavior for callers that omit them.
|
||||
max_seq_length: int = Field(0, ge = 0, le = 1048576)
|
||||
load_in_4bit: bool = Field(True)
|
||||
cache_type_kv: Optional[str] = Field(None)
|
||||
tensor_parallel: bool = Field(False)
|
||||
gpu_ids: Optional[List[int]] = Field(None)
|
||||
gpu_memory_mode: Literal["auto", "manual"] = Field(
|
||||
"auto",
|
||||
|
|
|
|||
|
|
@ -1004,8 +1004,13 @@ try:
|
|||
_DEFAULT_MAX_TOKENS_FLOOR,
|
||||
_DEFAULT_STREAM_STALL_TIMEOUT_S,
|
||||
_canonicalize_spec_mode,
|
||||
_extra_args_n_ubatch,
|
||||
_extra_args_set_spec_type,
|
||||
_hf_offline_if_dns_dead,
|
||||
_kv_bytes_per_elem,
|
||||
_kv_unified_from_args,
|
||||
_planned_main_cache_types,
|
||||
_swa_full_from_args_or_env,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import (
|
||||
|
|
@ -1043,8 +1048,13 @@ except ImportError:
|
|||
_DEFAULT_MAX_TOKENS_FLOOR,
|
||||
_DEFAULT_STREAM_STALL_TIMEOUT_S,
|
||||
_canonicalize_spec_mode,
|
||||
_extra_args_n_ubatch,
|
||||
_extra_args_set_spec_type,
|
||||
_hf_offline_if_dns_dead,
|
||||
_kv_bytes_per_elem,
|
||||
_kv_unified_from_args,
|
||||
_planned_main_cache_types,
|
||||
_swa_full_from_args_or_env,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import (
|
||||
|
|
@ -3349,6 +3359,10 @@ def _request_matches_loaded_settings(
|
|||
strip_offload = request.gpu_memory_mode == "manual",
|
||||
)
|
||||
)
|
||||
if not llama_backend.is_diffusion and llama_backend.swa_full != _swa_full_from_args_or_env(
|
||||
effective_extra
|
||||
):
|
||||
return False
|
||||
if not _tensor_parallel_matches_loaded(
|
||||
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
|
||||
):
|
||||
|
|
@ -4464,10 +4478,12 @@ def _estimate_gguf_kv_gb(
|
|||
max_seq_length: int,
|
||||
llama_extra_args: Optional[list[str]] = None,
|
||||
n_parallel: int = 1,
|
||||
cache_type_kv: Optional[str] = None,
|
||||
tensor_parallel: bool = False,
|
||||
) -> float:
|
||||
"""KV-cache VRAM (GB) at the larger of max_seq_length and any `--ctx-size`/`-c`
|
||||
override, over n_parallel slots, with the default f16 cache so the estimate is
|
||||
never below what the server allocates. 0 if metadata is unreadable."""
|
||||
override, over n_parallel slots, using the effective cache settings and managed
|
||||
launcher defaults. 0 if metadata is unreadable."""
|
||||
try:
|
||||
from core.inference.llama_server_args import parse_ctx_override
|
||||
|
||||
|
|
@ -4482,7 +4498,43 @@ def _estimate_gguf_kv_gb(
|
|||
ctx = max(max_seq_length or 0, ctx_override) or (probe._context_length or 0)
|
||||
if ctx <= 0:
|
||||
return 0.0
|
||||
kv = probe._estimate_kv_cache_bytes(ctx, n_parallel = max(1, n_parallel or 1))
|
||||
slots = max(1, n_parallel or 1)
|
||||
managed_kv_unified = bool(
|
||||
slots > 1
|
||||
and LlamaCppBackend.probe_server_capabilities().get("supports_kv_unified", False)
|
||||
)
|
||||
planned_cache_types = _planned_main_cache_types(
|
||||
cache_type_kv,
|
||||
llama_extra_args,
|
||||
)
|
||||
if tensor_parallel and any(
|
||||
cache_type not in LlamaCppBackend._TENSOR_PARALLEL_KV_TYPES
|
||||
for cache_type in planned_cache_types
|
||||
):
|
||||
# Tensor mode strips quantized axes, but a layer fallback restores
|
||||
# the original settings. Size for the larger successful outcome.
|
||||
tensor_cache_types = _planned_main_cache_types(None, None)
|
||||
cache_type_for_budget = max(
|
||||
(*planned_cache_types, *tensor_cache_types, "f16"),
|
||||
key = _kv_bytes_per_elem,
|
||||
)
|
||||
else:
|
||||
cache_type_for_budget = max(
|
||||
planned_cache_types,
|
||||
key = _kv_bytes_per_elem,
|
||||
)
|
||||
kv = probe._estimate_kv_cache_bytes(
|
||||
ctx,
|
||||
cache_type_for_budget,
|
||||
n_parallel = slots,
|
||||
swa_full = _swa_full_from_args_or_env(llama_extra_args),
|
||||
kv_unified = _kv_unified_from_args(
|
||||
llama_extra_args,
|
||||
default = managed_kv_unified,
|
||||
),
|
||||
n_ubatch = _extra_args_n_ubatch(llama_extra_args, n_ctx = ctx),
|
||||
flash_attn = False,
|
||||
)
|
||||
return kv / (1024**3)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not size GGUF KV cache for training guard: {e}")
|
||||
|
|
@ -4495,6 +4547,8 @@ def _estimate_gguf_required_gb(
|
|||
max_seq_length: int = 0,
|
||||
llama_extra_args: Optional[list[str]] = None,
|
||||
n_parallel: int = 1,
|
||||
cache_type_kv: Optional[str] = None,
|
||||
tensor_parallel: bool = False,
|
||||
) -> Optional[float]:
|
||||
"""Approximate GGUF VRAM (GB): quantized weights + companions, plus the KV
|
||||
cache for local files (unreadable pre-download for remote). None when nothing
|
||||
|
|
@ -4510,7 +4564,12 @@ def _estimate_gguf_required_gb(
|
|||
total_bytes += Path(f).stat().st_size
|
||||
if total_bytes > 0:
|
||||
return total_bytes / (1024**3) + _estimate_gguf_kv_gb(
|
||||
main, max_seq_length, llama_extra_args, n_parallel
|
||||
main,
|
||||
max_seq_length,
|
||||
llama_extra_args,
|
||||
n_parallel,
|
||||
cache_type_kv,
|
||||
tensor_parallel,
|
||||
)
|
||||
|
||||
repo = getattr(config, "gguf_hf_repo", None)
|
||||
|
|
@ -4651,6 +4710,8 @@ def _guard_chat_load_against_training(
|
|||
requested_gpu_ids: Optional[List[int]],
|
||||
llama_extra_args: Optional[list[str]] = None,
|
||||
n_parallel: int = 1,
|
||||
cache_type_kv: Optional[str] = None,
|
||||
tensor_parallel: bool = False,
|
||||
gpu_memory_mode: Literal["auto", "manual"] = "auto",
|
||||
) -> None:
|
||||
"""Protect active training from automatically placed chat-model loads.
|
||||
|
|
@ -4705,6 +4766,11 @@ def _guard_chat_load_against_training(
|
|||
max_seq_length = max_seq_length,
|
||||
llama_extra_args = llama_extra_args,
|
||||
n_parallel = n_parallel,
|
||||
cache_type_kv = cache_type_kv,
|
||||
tensor_parallel = (
|
||||
_effective_tensor_parallel(llama_extra_args, tensor_parallel)
|
||||
and (is_vulkan or LlamaCppBackend._effective_gpu_count(requested_gpu_ids) >= 2)
|
||||
),
|
||||
)
|
||||
if is_gguf
|
||||
else None
|
||||
|
|
@ -5460,6 +5526,8 @@ async def _load_model_impl(
|
|||
requested_gpu_ids = effective_gpu_ids,
|
||||
llama_extra_args = extra_llama_args,
|
||||
n_parallel = _n_parallel,
|
||||
cache_type_kv = request.cache_type_kv,
|
||||
tensor_parallel = bool(request.tensor_parallel),
|
||||
gpu_memory_mode = request.gpu_memory_mode,
|
||||
)
|
||||
|
||||
|
|
@ -6141,6 +6209,8 @@ async def validate_model(
|
|||
1,
|
||||
)
|
||||
),
|
||||
cache_type_kv = request.cache_type_kv,
|
||||
tensor_parallel = request.tensor_parallel,
|
||||
gpu_memory_mode = request.gpu_memory_mode,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -451,6 +451,9 @@ class TestChatLoadGuardRoute(unittest.TestCase):
|
|||
decision,
|
||||
gpu_memory_mode = "auto",
|
||||
requested_gpu_ids = None,
|
||||
llama_extra_args = None,
|
||||
cache_type_kv = None,
|
||||
tensor_parallel = False,
|
||||
):
|
||||
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
|
||||
with _stub_guard_deps(
|
||||
|
|
@ -463,6 +466,9 @@ class TestChatLoadGuardRoute(unittest.TestCase):
|
|||
load_in_4bit = True,
|
||||
max_seq_length = 0,
|
||||
requested_gpu_ids = requested_gpu_ids,
|
||||
llama_extra_args = llama_extra_args,
|
||||
cache_type_kv = cache_type_kv,
|
||||
tensor_parallel = tensor_parallel,
|
||||
gpu_memory_mode = gpu_memory_mode,
|
||||
)
|
||||
|
||||
|
|
@ -597,6 +603,32 @@ class TestChatLoadGuardRoute(unittest.TestCase):
|
|||
self.assertEqual(captured[0]["is_gguf"], True)
|
||||
self.assertEqual(captured[0]["required_override_gb"], 12.5)
|
||||
|
||||
def test_vulkan_gguf_estimate_keeps_tensor_cache_coercion(self):
|
||||
config = SimpleNamespace(is_gguf = True)
|
||||
estimate_kwargs = {}
|
||||
with (
|
||||
patch.object(
|
||||
self.route,
|
||||
"_estimate_gguf_required_gb",
|
||||
side_effect = lambda *args, **kwargs: estimate_kwargs.update(kwargs) or 12.5,
|
||||
),
|
||||
patch.object(
|
||||
self.route.LlamaCppBackend,
|
||||
"_effective_gpu_count",
|
||||
return_value = 0,
|
||||
),
|
||||
patch.object(self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True),
|
||||
):
|
||||
self._guard(
|
||||
config = config,
|
||||
training_active = True,
|
||||
decision = (True, {}),
|
||||
llama_extra_args = ["--split-mode", "tensor"],
|
||||
cache_type_kv = "q4_0",
|
||||
)
|
||||
self.assertEqual(estimate_kwargs["cache_type_kv"], "q4_0")
|
||||
self.assertTrue(estimate_kwargs["tensor_parallel"])
|
||||
|
||||
|
||||
class TestEffectiveLoadIn4bit(unittest.TestCase):
|
||||
@classmethod
|
||||
|
|
@ -745,7 +777,12 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
|
|||
# /load then 409s after the frontend has already unloaded.
|
||||
from models.inference import ValidateModelRequest
|
||||
|
||||
request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096)
|
||||
request = ValidateModelRequest(
|
||||
model_path = "unsloth/Qwen3-1.7B",
|
||||
max_seq_length = 4096,
|
||||
cache_type_kv = "f32",
|
||||
tensor_parallel = True,
|
||||
)
|
||||
cfg = SimpleNamespace(
|
||||
identifier = "unsloth/Qwen3-1.7B",
|
||||
display_name = "Qwen3-1.7B",
|
||||
|
|
@ -774,6 +811,8 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
|
|||
asyncio.run(self.route.validate_model(request, current_subject = "u"))
|
||||
self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"])
|
||||
self.assertIn("n_parallel", captured)
|
||||
self.assertEqual(captured.get("cache_type_kv"), "f32")
|
||||
self.assertTrue(captured.get("tensor_parallel"))
|
||||
|
||||
def test_metadata_probe_skips_training_guard(self):
|
||||
# A header-only probe (include_context_length) allocates no VRAM, so the
|
||||
|
|
@ -985,6 +1024,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
|
|||
|
||||
class _FakeBackend:
|
||||
_context_length = 2048
|
||||
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
|
||||
supports_kv_unified = True
|
||||
|
||||
def _read_gguf_metadata(self, path):
|
||||
pass
|
||||
|
|
@ -992,13 +1033,27 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
|
|||
def _can_estimate_kv(self):
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def probe_server_capabilities(cls):
|
||||
return {"supports_kv_unified": cls.supports_kv_unified}
|
||||
|
||||
def _estimate_kv_cache_bytes(
|
||||
self,
|
||||
ctx,
|
||||
cache_type = None,
|
||||
n_parallel = 1,
|
||||
swa_full = False,
|
||||
kv_unified = False,
|
||||
n_ubatch = None,
|
||||
flash_attn = True,
|
||||
):
|
||||
seen["ctx"] = ctx
|
||||
seen["cache_type"] = cache_type
|
||||
seen["n_parallel"] = n_parallel
|
||||
seen["swa_full"] = swa_full
|
||||
seen["kv_unified"] = kv_unified
|
||||
seen["n_ubatch"] = n_ubatch
|
||||
seen["flash_attn"] = flash_attn
|
||||
return ctx * n_parallel * (1024**2) # 1 MiB per ctx unit per slot
|
||||
|
||||
with patch.object(self.route, "LlamaCppBackend", _FakeBackend):
|
||||
|
|
@ -1009,6 +1064,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
|
|||
)
|
||||
self.assertEqual(seen["ctx"], 131072)
|
||||
self.assertEqual(seen["n_parallel"], 1) # default single slot
|
||||
self.assertFalse(seen["swa_full"])
|
||||
self.assertFalse(seen["flash_attn"])
|
||||
# override below max_seq_length -> larger (max_seq_length) wins
|
||||
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0)
|
||||
self.assertEqual(seen["ctx"], 4096)
|
||||
|
|
@ -1020,6 +1077,50 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
|
|||
# --parallel slots scale the cache the same way the launcher does
|
||||
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0)
|
||||
self.assertEqual(seen["n_parallel"], 4)
|
||||
self.assertTrue(seen["kv_unified"])
|
||||
# User extras are appended after Studio's managed default.
|
||||
r._estimate_gguf_kv_gb("m", 4096, ["--no-kv-unified"], 4)
|
||||
self.assertFalse(seen["kv_unified"])
|
||||
# An older binary without the flag keeps separate KV streams.
|
||||
_FakeBackend.supports_kv_unified = False
|
||||
r._estimate_gguf_kv_gb("m", 4096, None, 4)
|
||||
self.assertFalse(seen["kv_unified"])
|
||||
r._estimate_gguf_kv_gb("m", 4096, None, 1, "f32")
|
||||
self.assertEqual(seen["cache_type"], "f32")
|
||||
r._estimate_gguf_kv_gb("m", 4096, ["--cache-type-v", "f32"])
|
||||
self.assertEqual(seen["cache_type"], "f32")
|
||||
with patch.dict(self.route.os.environ, {"LLAMA_ARG_CACHE_TYPE_K": "f32"}):
|
||||
r._estimate_gguf_kv_gb("m", 4096)
|
||||
self.assertEqual(seen["cache_type"], "f32")
|
||||
with patch.dict(
|
||||
self.route.os.environ,
|
||||
{
|
||||
"LLAMA_ARG_CACHE_TYPE_K": "q4_0",
|
||||
"LLAMA_ARG_CACHE_TYPE_V": "q4_0",
|
||||
},
|
||||
):
|
||||
r._estimate_gguf_kv_gb("m", 4096)
|
||||
self.assertEqual(seen["cache_type"], "q4_0")
|
||||
r._estimate_gguf_kv_gb(
|
||||
"m",
|
||||
4096,
|
||||
["--cache-type-k", "q4_0", "--cache-type-v", "q4_0"],
|
||||
tensor_parallel = True,
|
||||
)
|
||||
self.assertEqual(seen["cache_type"], "f16")
|
||||
r._estimate_gguf_kv_gb(
|
||||
"m",
|
||||
4096,
|
||||
["--cache-type-k", "f32", "--cache-type-v", "q4_0"],
|
||||
tensor_parallel = True,
|
||||
)
|
||||
self.assertEqual(seen["cache_type"], "f32")
|
||||
# Full SWA mode follows the same pass-through args as the launcher.
|
||||
r._estimate_gguf_kv_gb("m", 4096, ["--swa_full"])
|
||||
self.assertTrue(seen["swa_full"])
|
||||
r._estimate_gguf_kv_gb("m", 4096, ["--kv_unified", "--ubatch_size", "256"])
|
||||
self.assertTrue(seen["kv_unified"])
|
||||
self.assertEqual(seen["n_ubatch"], 256)
|
||||
|
||||
|
||||
# ── load_model integration: authoritative 409, and no unload before refusal ──
|
||||
|
|
|
|||
|
|
@ -183,11 +183,12 @@ def test_already_in_target_state_reloads_on_mode_change(loaded, requested):
|
|||
assert _target_state(_loaded_backend(loaded), requested) is False
|
||||
|
||||
|
||||
def test_already_in_target_state_ignores_mode_for_diffusion():
|
||||
def test_already_in_target_state_ignores_mode_for_diffusion(monkeypatch):
|
||||
# The diffusion runner is mode-agnostic (always "auto"), so a standing manual
|
||||
# preference must not force a needless reload.
|
||||
backend = _loaded_backend("auto")
|
||||
backend._is_diffusion = True
|
||||
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
|
||||
assert _target_state(backend, "manual") is True
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -76,6 +76,39 @@ from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
|
|||
# Helpers
|
||||
|
||||
|
||||
def _runtime_kv_cells(
|
||||
n_ctx: int,
|
||||
*,
|
||||
slots: int = 1,
|
||||
unified: bool = True,
|
||||
) -> int:
|
||||
"""Total KV cells allocated by llama.cpp across all streams."""
|
||||
slots = max(1, slots)
|
||||
padded_ctx = ((n_ctx + 255) // 256) * 256
|
||||
streams = 1 if unified else slots
|
||||
cells_per_stream = padded_ctx if unified else ((max(1, padded_ctx // slots) + 255) // 256) * 256
|
||||
return cells_per_stream * streams
|
||||
|
||||
|
||||
def _runtime_swa_cells(
|
||||
n_ctx: int,
|
||||
sliding_window: int,
|
||||
*,
|
||||
slots: int = 1,
|
||||
unified: bool = True,
|
||||
n_ubatch: int = 512,
|
||||
) -> tuple[int, int]:
|
||||
"""Return total non-SWA and compact-SWA cells allocated by llama.cpp."""
|
||||
slots = max(1, slots)
|
||||
streams = 1 if unified else slots
|
||||
base_cells = _runtime_kv_cells(n_ctx, slots = slots, unified = unified)
|
||||
cells_per_stream = base_cells // streams
|
||||
swa_limit = sliding_window * (slots if unified else 1) + n_ubatch
|
||||
swa_cells_per_stream = min(cells_per_stream, swa_limit)
|
||||
swa_cells_per_stream = ((swa_cells_per_stream + 255) // 256) * 256
|
||||
return base_cells, swa_cells_per_stream * streams
|
||||
|
||||
|
||||
def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
|
||||
"""Build a minimal GGUF v3 blob with the given KV metadata.
|
||||
|
||||
|
|
@ -789,7 +822,7 @@ class TestMLAEstimation:
|
|||
b = self._mla_backend()
|
||||
result = b._estimate_kv_cache_bytes(1000, "f16")
|
||||
# n_layers * ctx * 1 * key_len(576) * 2
|
||||
expected = 61 * 1000 * 1 * 576 * 2
|
||||
expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2
|
||||
assert result == expected
|
||||
|
||||
def test_mla_fallback_when_no_key_length(self):
|
||||
|
|
@ -797,14 +830,14 @@ class TestMLAEstimation:
|
|||
b = self._mla_backend(_kv_key_length = None)
|
||||
# default _key_length_mla=192, so rope_dim=192
|
||||
result = b._estimate_kv_cache_bytes(1000, "f16")
|
||||
expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704
|
||||
expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 192) * 2 # 704
|
||||
assert result == expected
|
||||
|
||||
def test_mla_fallback_no_key_length_mla(self):
|
||||
"""No key_length and no key_length_mla: fall back to +64."""
|
||||
b = self._mla_backend(_kv_key_length = None, _key_length_mla = None)
|
||||
result = b._estimate_kv_cache_bytes(1000, "f16")
|
||||
expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576
|
||||
expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 64) * 2 # 576
|
||||
assert result == expected
|
||||
|
||||
def test_mla_defaults_n_kv_to_1_when_heads_absent(self):
|
||||
|
|
@ -812,7 +845,7 @@ class TestMLAEstimation:
|
|||
b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set
|
||||
result = b._estimate_kv_cache_bytes(1000, "f16")
|
||||
# Uses n_kv_mla=1, NOT n_heads=128
|
||||
expected = 61 * 1000 * 1 * 576 * 2
|
||||
expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2
|
||||
assert result == expected
|
||||
|
||||
def test_mla_q4_quantization(self):
|
||||
|
|
@ -821,7 +854,7 @@ class TestMLAEstimation:
|
|||
result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0")
|
||||
assert result_q4 < result_f16
|
||||
# q4_0 bpe = 0.5625, f16 bpe = 2.0
|
||||
assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625)
|
||||
assert result_q4 == int(61 * _runtime_kv_cells(1000) * 1 * 576 * 0.5625)
|
||||
|
||||
|
||||
# D. Path 2: Hybrid Mamba Estimation
|
||||
|
|
@ -910,9 +943,8 @@ class TestSlidingWindowEstimation:
|
|||
n_global = max(1, 62 // 4) # 15
|
||||
n_swa = 62 - n_global # 47
|
||||
kv_per = 16 * (128 + 128) * 2
|
||||
# SWA cache is double-buffered: 2 * sliding_window cells, capped at n_ctx.
|
||||
swa_cells = min(131072, 2 * 1024)
|
||||
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
||||
base_cells, swa_cells = _runtime_swa_cells(131072, 1024)
|
||||
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
|
||||
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
||||
|
||||
def test_gpt_oss(self):
|
||||
|
|
@ -929,8 +961,8 @@ class TestSlidingWindowEstimation:
|
|||
n_global = max(1, 24 // 4) # 6
|
||||
n_swa = 24 - n_global # 18
|
||||
kv_per = 8 * (64 + 64) * 2
|
||||
swa_cells = min(131072, 2 * 128)
|
||||
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
||||
base_cells, swa_cells = _runtime_swa_cells(131072, 128)
|
||||
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
|
||||
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
||||
|
||||
def test_gemma4_per_layer_swa_metadata(self):
|
||||
|
|
@ -952,21 +984,67 @@ class TestSlidingWindowEstimation:
|
|||
sliding_layers = 25
|
||||
|
||||
def expected(ctx):
|
||||
full = full_layers * ctx * 2 * (512 + 512) * 2
|
||||
sliding = sliding_layers * min(ctx, 2 * 1024) * 8 * (256 + 256) * 2
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
|
||||
full = full_layers * base_cells * 2 * (512 + 512) * 2
|
||||
sliding = sliding_layers * swa_cells * 8 * (256 + 256) * 2
|
||||
return int(full + sliding)
|
||||
|
||||
for ctx in (4096, 46500, 262144):
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx)
|
||||
|
||||
def test_gemma4_flash_attn_off_pads_v_to_model_max(self):
|
||||
b = self._swa_backend(
|
||||
_n_layers = 35,
|
||||
_n_kv_heads = 1,
|
||||
_n_heads = 8,
|
||||
_embedding_length = 1536,
|
||||
_kv_key_length = 512,
|
||||
_kv_value_length = 512,
|
||||
_sliding_window = 512,
|
||||
_sliding_window_pattern = [True, True, True, True, False] * 7,
|
||||
_kv_key_length_swa = 256,
|
||||
_kv_value_length_swa = 256,
|
||||
_shared_kv_layers = 20,
|
||||
)
|
||||
ctx = 5000
|
||||
slots = 3
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, 512, slots = slots, unified = True)
|
||||
max_v_width = 512
|
||||
expected = (
|
||||
3 * base_cells * (512 + max_v_width) * 2 + 12 * swa_cells * (256 + max_v_width) * 2
|
||||
)
|
||||
actual = b._estimate_kv_cache_bytes(
|
||||
ctx,
|
||||
"f16",
|
||||
n_parallel = slots,
|
||||
flash_attn = False,
|
||||
)
|
||||
assert actual == expected
|
||||
assert actual == 66 * 1024**2
|
||||
assert actual > b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
|
||||
|
||||
def test_flash_attn_off_prices_quantized_v_retry_as_f16(self):
|
||||
b = self._swa_backend(
|
||||
_n_layers = 2,
|
||||
_n_kv_heads = None,
|
||||
_n_kv_heads_by_layer = [8, 2],
|
||||
_sliding_window_pattern = [True, False],
|
||||
_kv_key_length_swa = 64,
|
||||
_kv_value_length_swa = 64,
|
||||
)
|
||||
off = b._estimate_kv_cache_bytes(4096, "q4_0", flash_attn = False)
|
||||
on = b._estimate_kv_cache_bytes(4096, "q4_0")
|
||||
assert off > on
|
||||
|
||||
def test_ctx_smaller_than_window(self):
|
||||
"""When ctx < 2 * sliding_window, SWA cache caps at ctx."""
|
||||
"""When context is smaller than the compact allowance, SWA caps at context."""
|
||||
b = self._swa_backend(_sliding_window = 8192)
|
||||
n_global = max(1, 62 // 4) # 15
|
||||
n_swa = 62 - n_global # 47
|
||||
kv_per = 16 * (128 + 128) * 2
|
||||
ctx = 4096
|
||||
expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 2 * 8192) * kv_per)
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, 8192)
|
||||
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
||||
|
||||
def test_odd_layer_count(self):
|
||||
|
|
@ -974,7 +1052,8 @@ class TestSlidingWindowEstimation:
|
|||
n_global = max(1, 63 // 4) # 15
|
||||
n_swa = 63 - n_global # 48
|
||||
kv_per = 16 * (128 + 128) * 2
|
||||
expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per)
|
||||
base_cells, swa_cells = _runtime_swa_cells(1000, 1024)
|
||||
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
|
||||
assert b._estimate_kv_cache_bytes(1000, "f16") == expected
|
||||
|
||||
|
||||
|
|
@ -1086,8 +1165,7 @@ class TestPathPriority:
|
|||
b._full_attention_interval = 4
|
||||
b._sliding_window = 1024 # Would trigger SWA
|
||||
|
||||
# MLA: 61 * 1000 * 1 * 576 * 2
|
||||
expected_mla = int(61 * 1000 * 1 * 576 * 2)
|
||||
expected_mla = int(61 * _runtime_kv_cells(1000) * 1 * 576 * 2)
|
||||
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla
|
||||
|
||||
def test_hybrid_over_swa(self):
|
||||
|
|
@ -1104,7 +1182,7 @@ class TestPathPriority:
|
|||
b._sliding_window = 1024 # Would trigger SWA
|
||||
|
||||
n_attn = 64 // 4
|
||||
expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2)
|
||||
expected_hybrid = int(n_attn * _runtime_kv_cells(1000) * 4 * (256 + 256) * 2)
|
||||
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid
|
||||
|
||||
def test_all_paths_produce_different_values(self):
|
||||
|
|
@ -1192,7 +1270,7 @@ class TestQuantization:
|
|||
b._kv_key_length = 64
|
||||
b._kv_value_length = 64
|
||||
result = b._estimate_kv_cache_bytes(1000, cache_type)
|
||||
expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe)
|
||||
expected = int(10 * _runtime_kv_cells(1000) * 1 * (64 + 64) * expected_bpe)
|
||||
assert result == expected
|
||||
|
||||
|
||||
|
|
@ -1221,7 +1299,7 @@ class TestEdgeCases:
|
|||
b._kv_key_length = 64
|
||||
b._kv_value_length = 64
|
||||
result = b._estimate_kv_cache_bytes(1, "f16")
|
||||
assert result == int(10 * 1 * 1 * (64 + 64) * 2)
|
||||
assert result == int(10 * _runtime_kv_cells(1) * 1 * (64 + 64) * 2)
|
||||
|
||||
def test_very_large_context(self):
|
||||
"""1M context should not overflow or crash."""
|
||||
|
|
@ -1242,7 +1320,7 @@ class TestEdgeCases:
|
|||
b._kv_key_length = 64
|
||||
b._kv_value_length = 64
|
||||
result = b._estimate_kv_cache_bytes(100, "f16")
|
||||
expected = int(10 * 100 * 8 * (64 + 64) * 2)
|
||||
expected = int(10 * _runtime_kv_cells(100) * 8 * (64 + 64) * 2)
|
||||
assert result == expected
|
||||
|
||||
def test_both_heads_none_falls_to_one(self):
|
||||
|
|
@ -1253,7 +1331,7 @@ class TestEdgeCases:
|
|||
b._kv_key_length = 64
|
||||
b._kv_value_length = 64
|
||||
result = b._estimate_kv_cache_bytes(100, "f16")
|
||||
expected = int(10 * 100 * 1 * (64 + 64) * 2)
|
||||
expected = int(10 * _runtime_kv_cells(100) * 1 * (64 + 64) * 2)
|
||||
assert result == expected
|
||||
|
||||
|
||||
|
|
@ -1335,12 +1413,21 @@ class TestServerFlags:
|
|||
assert with_cp_full == no_cp_full
|
||||
assert with_cp > b._estimate_kv_cache_bytes(8192, "f16")
|
||||
|
||||
def test_compact_swa_includes_ubatch_headroom_and_padding(self):
|
||||
b = self._swa_backend(_sliding_window = 128)
|
||||
ctx = 8192
|
||||
result = b._estimate_kv_cache_bytes(ctx, "f16", n_ubatch = 512)
|
||||
per_token = 4 * (256 + 256) * 2
|
||||
n_swa = sum(b._sliding_window_pattern)
|
||||
n_global = b._n_layers - n_swa
|
||||
expected = n_global * ctx * per_token + n_swa * 768 * per_token
|
||||
assert result == expected
|
||||
|
||||
# ── --parallel + --kv-unified ──────────────────────────────────
|
||||
# Verified against llama-server: non-SWA caches partition n_ctx across
|
||||
# slots (total memory constant); only SWA layers scale with --parallel.
|
||||
# --kv-unified is a no-op for memory math (kept for API forward-compat).
|
||||
# non-unified streams. Compact SWA sizing depends on the stream layout.
|
||||
|
||||
def test_gqa_kv_constant_across_parallel(self):
|
||||
def test_gqa_kv_constant_for_aligned_stream_divisions(self):
|
||||
b = self._gqa_backend()
|
||||
baseline = b._estimate_kv_cache_bytes(4096, "f16")
|
||||
for slots in (1, 2, 4, 8):
|
||||
|
|
@ -1359,7 +1446,7 @@ class TestServerFlags:
|
|||
== baseline
|
||||
)
|
||||
|
||||
def test_swa_path_scales_only_swa_portion(self):
|
||||
def test_swa_path_matches_aligned_stream_layout(self):
|
||||
b = self._swa_backend()
|
||||
ctx = 8192
|
||||
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
||||
|
|
@ -1367,27 +1454,27 @@ class TestServerFlags:
|
|||
swa = b._sliding_window
|
||||
per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
||||
per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back
|
||||
per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa)
|
||||
global_bytes = sum(
|
||||
ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
|
||||
base_cells * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
|
||||
)
|
||||
swa_bytes_per_slot = sum(
|
||||
per_slot_swa_cells * per_token_swa
|
||||
for f in b._sliding_window_pattern[: b._n_layers]
|
||||
if f
|
||||
swa_bytes = sum(
|
||||
swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
|
||||
)
|
||||
# Sanity: parallel=1 reproduces baseline exactly
|
||||
assert global_bytes + swa_bytes_per_slot == baseline
|
||||
# Only the SWA portion scales by parallel
|
||||
assert global_bytes + swa_bytes == baseline
|
||||
for slots in (1, 2, 3, 4):
|
||||
scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
||||
# SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa
|
||||
per_slot_ctx = max(1, ctx // slots)
|
||||
cells = min(ctx, 2 * swa, per_slot_ctx)
|
||||
swa_bps = sum(
|
||||
cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
|
||||
expected_global = sum(
|
||||
base_cells * per_token_global
|
||||
for f in b._sliding_window_pattern[: b._n_layers]
|
||||
if not f
|
||||
)
|
||||
assert scaled == global_bytes + slots * swa_bps
|
||||
expected_swa = sum(
|
||||
swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
|
||||
)
|
||||
assert scaled == expected_global + expected_swa
|
||||
|
||||
def test_mla_kv_constant_across_parallel(self):
|
||||
b = LlamaCppBackend()
|
||||
|
|
@ -1444,19 +1531,17 @@ class TestServerFlags:
|
|||
ctx = 8192
|
||||
swa = b._sliding_window
|
||||
per_token = 4 * (256 + 256) * 2
|
||||
global_bytes = sum(
|
||||
ctx * per_token for f in b._sliding_window_pattern[: b._n_layers] if not f
|
||||
)
|
||||
n_swa_layers = sum(1 for f in b._sliding_window_pattern[: b._n_layers] if f)
|
||||
slots = 3
|
||||
per_slot_ctx = max(1, ctx // slots)
|
||||
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
||||
swa_bytes_per_slot = n_swa_layers * swa_cells * per_token
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
|
||||
n_global_layers = b._n_layers - n_swa_layers
|
||||
global_bytes = n_global_layers * base_cells * per_token
|
||||
swa_bytes = n_swa_layers * swa_cells * per_token
|
||||
cp_extra_per_slot = n_swa_layers * 4 * swa * per_token # 4 checkpoints
|
||||
flagged = b._estimate_kv_cache_bytes(
|
||||
ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False
|
||||
)
|
||||
assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot)
|
||||
assert flagged == global_bytes + swa_bytes + slots * cp_extra_per_slot
|
||||
|
||||
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
|
||||
|
||||
|
|
@ -1535,22 +1620,40 @@ class TestServerFlags:
|
|||
assert fitted_default == ctx
|
||||
assert fitted_full < ctx
|
||||
|
||||
def test_tensor_planner_threads_swa_full_through_estimator(self):
|
||||
b = self._swa_backend()
|
||||
estimate = b._estimate_kv_cache_bytes
|
||||
calls = []
|
||||
|
||||
def record(*args, **kwargs):
|
||||
calls.append(kwargs)
|
||||
return estimate(*args, **kwargs)
|
||||
|
||||
b._estimate_kv_cache_bytes = record
|
||||
b._plan_tensor_parallel(
|
||||
[(0, 32768), (1, 32768)],
|
||||
1024**3,
|
||||
8192,
|
||||
cache_type_kv = "f16",
|
||||
swa_full = True,
|
||||
flash_attn = False,
|
||||
)
|
||||
assert calls
|
||||
assert all(call["swa_full"] is True for call in calls)
|
||||
assert all(call["flash_attn"] is False for call in calls)
|
||||
|
||||
|
||||
# J2.5. --parallel N memory accounting (per-layer-type scaling rule)
|
||||
|
||||
|
||||
class TestParallelSWAScaling:
|
||||
"""Per-layer-type scaling rule vs the closed form measured from
|
||||
llama-server. Empirical formula on Gemma-3 270m at ctx=8192:
|
||||
total_kv = 24 + parallel * 15 (MiB).
|
||||
"""Per-layer-type scaling rule measured from llama-server.
|
||||
|
||||
Rule (verified vs ``llama-server`` log on real GGUFs):
|
||||
* non-SWA layers: total cells = n_ctx, partitioned across slots,
|
||||
memory CONSTANT in n_parallel.
|
||||
* SWA layers: per-slot cells = 2 * sliding_window (clamped at
|
||||
n_ctx and at per_slot_ctx); memory LINEAR in n_parallel.
|
||||
* --kv-unified is a no-op for memory math; both modes give the
|
||||
same total in measured cases.
|
||||
* non-SWA layers use the padded per-stream context.
|
||||
* compact SWA adds ubatch headroom and pads to 256 cells.
|
||||
* unified mode uses one stream with all slot windows.
|
||||
* non-unified mode allocates one stream per slot.
|
||||
"""
|
||||
|
||||
def _gqa_backend(self, **overrides):
|
||||
|
|
@ -1586,7 +1689,7 @@ class TestParallelSWAScaling:
|
|||
setattr(b, k, v)
|
||||
return b
|
||||
|
||||
# ── non-SWA paths: constant ────────────────────────────────────
|
||||
# ── non-SWA paths: constant when stream divisions are aligned ──
|
||||
|
||||
def test_pure_gqa_constant_across_parallel(self):
|
||||
b = self._gqa_backend()
|
||||
|
|
@ -1633,25 +1736,53 @@ class TestParallelSWAScaling:
|
|||
for slots in (1, 2, 4, 8):
|
||||
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
||||
|
||||
# ── SWA paths: scale only the SWA portion ──────────────────────
|
||||
def test_non_swa_paths_follow_unaligned_stream_padding(self):
|
||||
mla = LlamaCppBackend()
|
||||
mla._n_layers = 60
|
||||
mla._n_kv_heads = 1
|
||||
mla._kv_lora_rank = 512
|
||||
mla._key_length_mla = 64
|
||||
mla._kv_key_length = 576
|
||||
|
||||
def test_swa_pattern_scales_only_swa_portion(self):
|
||||
hybrid = LlamaCppBackend()
|
||||
hybrid._n_layers = 64
|
||||
hybrid._n_kv_heads = 16
|
||||
hybrid._n_heads = 32
|
||||
hybrid._embedding_length = 4096
|
||||
hybrid._kv_key_length = 128
|
||||
hybrid._kv_value_length = 128
|
||||
hybrid._ssm_inner_size = 4096
|
||||
hybrid._full_attention_interval = 4
|
||||
|
||||
legacy = LlamaCppBackend()
|
||||
legacy._n_layers = 32
|
||||
legacy._n_kv_heads = 8
|
||||
legacy._n_heads = 8
|
||||
legacy._embedding_length = 4096
|
||||
|
||||
for backend in (self._gqa_backend(), mla, hybrid, legacy):
|
||||
bytes_per_cell = backend._estimate_kv_cache_bytes(256, "f16") // 256
|
||||
unified = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = True)
|
||||
separate = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = False)
|
||||
assert unified == 5120 * bytes_per_cell
|
||||
assert separate == 5376 * bytes_per_cell
|
||||
|
||||
# ── SWA paths: aligned stream scaling ──────────────────────────
|
||||
|
||||
def test_swa_pattern_matches_aligned_stream_layout(self):
|
||||
b = self._swa_backend()
|
||||
ctx = 8192
|
||||
swa = b._sliding_window
|
||||
per_token = 1 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
||||
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
||||
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
||||
global_bytes = n_global * ctx * per_token
|
||||
for slots in (1, 2, 4, 8):
|
||||
per_slot_ctx = max(1, ctx // slots)
|
||||
cells = min(ctx, 2 * swa, per_slot_ctx)
|
||||
swa_bps = n_swa * cells * per_token
|
||||
for unified in (True, False):
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified)
|
||||
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
|
||||
assert got == global_bytes + slots * swa_bps
|
||||
assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token)
|
||||
|
||||
def test_swa_fallback_scales_only_swa_portion(self):
|
||||
def test_swa_fallback_matches_aligned_stream_layout(self):
|
||||
# No per-layer pattern -> 1/4-global heuristic.
|
||||
b = self._swa_backend(_sliding_window_pattern = None)
|
||||
ctx = 8192
|
||||
|
|
@ -1660,34 +1791,28 @@ class TestParallelSWAScaling:
|
|||
n_global = max(1, n_layers // 4)
|
||||
n_swa = n_layers - n_global
|
||||
per_token = 1 * (256 + 256) * 2
|
||||
global_bytes = n_global * ctx * per_token
|
||||
for slots in (1, 2, 4, 8):
|
||||
per_slot_ctx = max(1, ctx // slots)
|
||||
cells = min(ctx, 2 * swa, per_slot_ctx)
|
||||
swa_bps = n_swa * cells * per_token
|
||||
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
|
||||
assert got == global_bytes + slots * swa_bps
|
||||
for unified in (True, False):
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified)
|
||||
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
|
||||
assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token)
|
||||
|
||||
def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self):
|
||||
# ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024.
|
||||
# SWA cells clamp at per_slot_ctx (512), not 2*sliding.
|
||||
# ctx=4096 / slots=8 gives a 512-cell stream, which caps compact SWA.
|
||||
b = self._swa_backend()
|
||||
ctx = 4096
|
||||
per_slot_ctx_at_8 = ctx // 8
|
||||
assert per_slot_ctx_at_8 < 2 * b._sliding_window
|
||||
# Build expected with the clamped formula
|
||||
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
||||
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
||||
per_token = 1 * (256 + 256) * 2
|
||||
global_bytes = n_global * ctx * per_token
|
||||
cells = min(ctx, 2 * b._sliding_window, per_slot_ctx_at_8)
|
||||
assert cells == per_slot_ctx_at_8
|
||||
expected = global_bytes + 8 * (n_swa * cells * per_token)
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, b._sliding_window, slots = 8, unified = False)
|
||||
assert swa_cells == 8 * per_slot_ctx_at_8
|
||||
expected = n_global * base_cells * per_token + n_swa * swa_cells * per_token
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8, kv_unified = False) == expected
|
||||
|
||||
def test_swa_full_does_not_scale_under_parallel(self):
|
||||
# swa_full forces every layer to n_ctx -> all-global GQA-style
|
||||
# total, constant in parallel.
|
||||
def test_swa_full_constant_for_aligned_stream_divisions(self):
|
||||
# swa_full forces every layer to n_ctx. This aligned context remains
|
||||
# constant across the tested stream divisions.
|
||||
b = self._swa_backend()
|
||||
ctx = 8192
|
||||
baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
||||
|
|
@ -1696,25 +1821,32 @@ class TestParallelSWAScaling:
|
|||
b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline
|
||||
)
|
||||
|
||||
# ── kv_unified: no-op for memory math ──────────────────────────
|
||||
# ── kv_unified stream layout ────────────────────────────────────
|
||||
|
||||
def test_kv_unified_is_no_op_for_memory_math(self):
|
||||
# unified=True and unified=False must give the same total bytes
|
||||
# for every backend type and parallel value.
|
||||
backends = [
|
||||
("gqa", self._gqa_backend()),
|
||||
("swa", self._swa_backend()),
|
||||
]
|
||||
for label, b in backends:
|
||||
for slots in (1, 2, 4, 8):
|
||||
u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
|
||||
nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
|
||||
assert u == nu, f"{label} parallel={slots} unified-mismatch"
|
||||
def test_kv_unified_changes_only_compact_swa_for_aligned_context(self):
|
||||
gqa = self._gqa_backend()
|
||||
swa = self._swa_backend()
|
||||
for slots in (1, 2, 4, 8):
|
||||
gqa_unified = gqa._estimate_kv_cache_bytes(
|
||||
8192, "f16", n_parallel = slots, kv_unified = True
|
||||
)
|
||||
gqa_separate = gqa._estimate_kv_cache_bytes(
|
||||
8192, "f16", n_parallel = slots, kv_unified = False
|
||||
)
|
||||
assert gqa_unified == gqa_separate
|
||||
|
||||
swa_unified = swa._estimate_kv_cache_bytes(
|
||||
8192, "f16", n_parallel = slots, kv_unified = True
|
||||
)
|
||||
swa_separate = swa._estimate_kv_cache_bytes(
|
||||
8192, "f16", n_parallel = slots, kv_unified = False
|
||||
)
|
||||
assert (swa_unified == swa_separate) is (slots == 1)
|
||||
|
||||
# ── Empirical Gemma-3 270m formula ─────────────────────────────
|
||||
|
||||
def test_matches_empirical_gemma3_270m_formula(self):
|
||||
"""Exact match against the formula measured from llama-server:
|
||||
"""Exact match against the non-unified formula measured from llama-server:
|
||||
total_kv = 24 + parallel * 15 (MiB) at ctx=8192.
|
||||
|
||||
Geometry: 18 layers (3 global + 15 SWA), n_kv=1, head_dim=256,
|
||||
|
|
@ -1736,12 +1868,16 @@ class TestParallelSWAScaling:
|
|||
# Confirm pattern shape
|
||||
assert sum(b._sliding_window_pattern) == n_swa
|
||||
for slots, expected_mib in [(1, 39), (2, 54), (4, 84)]:
|
||||
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots)
|
||||
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
|
||||
got_mib = got_bytes / (1024 * 1024)
|
||||
assert (
|
||||
got_mib == expected_mib
|
||||
), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB"
|
||||
|
||||
for slots, expected_mib in [(1, 39), (2, 46.5), (4, 61.5)]:
|
||||
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
|
||||
assert got_bytes / (1024 * 1024) == expected_mib
|
||||
|
||||
|
||||
# J3. shared_kv_layers (Gemma 3n / Gemma 4)
|
||||
|
||||
|
|
@ -1844,8 +1980,8 @@ class TestSharedKVLayers:
|
|||
assert sliding_in_unshared == 16
|
||||
assert full_in_unshared == 4
|
||||
kv_per = 4 * (256 + 256) * 2
|
||||
swa_cells = min(ctx, 2 * 1024)
|
||||
expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
|
||||
expected = full_in_unshared * base_cells * kv_per + sliding_in_unshared * swa_cells * kv_per
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
||||
|
||||
def test_shared_layers_reduces_estimate(self):
|
||||
|
|
@ -1875,8 +2011,8 @@ class TestSharedKVLayers:
|
|||
n_global = max(1, n_layers_kv // 4) # 5
|
||||
n_swa = n_layers_kv - n_global # 15
|
||||
kv_per = 4 * (256 + 256) * 2
|
||||
swa_cells = min(ctx, 2 * 1024)
|
||||
expected = n_global * ctx * kv_per + n_swa * swa_cells * kv_per
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
|
||||
expected = n_global * base_cells * kv_per + n_swa * swa_cells * kv_per
|
||||
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
||||
|
||||
def test_shared_floors_at_one_layer(self):
|
||||
|
|
@ -1896,13 +2032,12 @@ class TestSharedKVLayers:
|
|||
unshared_pattern = b._sliding_window_pattern[:20] # 35 - 15 shared
|
||||
sliding_in_unshared = sum(unshared_pattern)
|
||||
global_in_unshared = len(unshared_pattern) - sliding_in_unshared
|
||||
global_bytes = global_in_unshared * ctx * per_token
|
||||
slots = 3
|
||||
per_slot_ctx = max(1, ctx // slots)
|
||||
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
||||
swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token
|
||||
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
|
||||
global_bytes = global_in_unshared * base_cells * per_token
|
||||
swa_bytes = sliding_in_unshared * swa_cells * per_token
|
||||
flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
||||
assert flagged == global_bytes + slots * swa_bytes_per_slot
|
||||
assert flagged == global_bytes + swa_bytes
|
||||
|
||||
def test_composes_with_ctx_checkpoints(self):
|
||||
b = self._gemma3n_backend()
|
||||
|
|
@ -2036,14 +2171,14 @@ class TestLifecycle:
|
|||
)
|
||||
assert b._can_estimate_kv()
|
||||
result = b._estimate_kv_cache_bytes(131072, "f16")
|
||||
# gemma3 -> period 6 from bootstrap; SWA cache double-buffered to
|
||||
# 2 * sliding_window cells.
|
||||
# gemma3 uses period 6 from the bootstrap resolver.
|
||||
period = 6
|
||||
kv_per = 16 * 256 * 2
|
||||
base_cells, swa_cells = _runtime_swa_cells(131072, 1024)
|
||||
expected = 0
|
||||
for i in range(62):
|
||||
is_swa = (i + 1) % period != 0
|
||||
layer_ctx = min(131072, 2 * 1024) if is_swa else 131072
|
||||
layer_ctx = swa_cells if is_swa else base_cells
|
||||
expected += layer_ctx * kv_per
|
||||
assert result == expected
|
||||
|
||||
|
|
|
|||
|
|
@ -221,6 +221,18 @@ class TestFlashAttnOff:
|
|||
assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"]
|
||||
assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"]
|
||||
|
||||
@pytest.mark.parametrize("value", ["on", "enabled", "true", "1", "auto", "-1"])
|
||||
def test_flips_every_enabled_value(self, value):
|
||||
assert _flash_off(["llama-server", "--flash-attn", value]) == [
|
||||
"llama-server",
|
||||
"--flash-attn",
|
||||
"off",
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("value", ["off", "disabled", "false", "0"])
|
||||
def test_none_for_every_disabled_value(self, value):
|
||||
assert _flash_off(["llama-server", "--flash-attn", value]) is None
|
||||
|
||||
def test_flips_every_occurrence_last_wins(self):
|
||||
# extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins,
|
||||
# so one leftover 'on' would re-crash the retry. Every enable must flip.
|
||||
|
|
@ -384,6 +396,10 @@ class TestFlashAttnOffQuantizedKvCache:
|
|||
out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"])
|
||||
assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"]
|
||||
|
||||
def test_underscore_alias_flash_attn_is_disabled(self):
|
||||
out = _flash_off(["llama-server", "--flash_attn=on"])
|
||||
assert out == ["llama-server", "--flash_attn=off"]
|
||||
|
||||
def test_underscore_value_not_normalized_for_nonquantized(self):
|
||||
# Only the flag name is canonicalized; a non-quantized type value is
|
||||
# matched verbatim and left untouched (no spurious reset).
|
||||
|
|
|
|||
|
|
@ -63,7 +63,9 @@ from core.inference.llama_cpp import (
|
|||
_extra_args_set_any_flag,
|
||||
_extra_args_set_spec_type,
|
||||
_is_mtp_model_name,
|
||||
_kv_unified_from_args,
|
||||
_mla_mtp_auto_enabled,
|
||||
_swa_full_from_args_or_env,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -147,6 +149,41 @@ def test_is_mtp_model_name_handles_none():
|
|||
assert _is_mtp_model_name("", "") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flag", ["--swa-full", "--swa_full"])
|
||||
def test_swa_full_detects_llama_cpp_long_flag_spellings(flag):
|
||||
assert _swa_full_from_args_or_env([flag], {}) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["on", "enabled", "true", "1"])
|
||||
def test_swa_full_detects_llama_cpp_env_truth_values(value):
|
||||
assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["", "off", "yes", "TRUE", " true ", "0"])
|
||||
def test_swa_full_rejects_values_llama_cpp_treats_as_false(value):
|
||||
assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is False
|
||||
|
||||
|
||||
def test_swa_full_cli_wins_when_env_is_false():
|
||||
assert _swa_full_from_args_or_env(["--swa-full"], {"LLAMA_ARG_SWA_FULL": "0"}) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flag", ["--kv-unified", "--kv_unified", "-kvu"])
|
||||
def test_kv_unified_detects_enable_aliases(flag):
|
||||
assert _kv_unified_from_args([flag]) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flag", ["--no-kv-unified", "--no_kv_unified", "-no-kvu"])
|
||||
def test_kv_unified_detects_disable_aliases(flag):
|
||||
assert _kv_unified_from_args(["--kv-unified", flag]) is False
|
||||
|
||||
|
||||
def test_kv_unified_uses_environment_before_cli():
|
||||
assert _kv_unified_from_args([], env = {"LLAMA_ARG_KV_UNIFIED": "true"}) is True
|
||||
assert _kv_unified_from_args([], default = True, env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True
|
||||
assert _kv_unified_from_args(["--kv-unified"], env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True
|
||||
|
||||
|
||||
def test_is_mtp_model_name_detects_marker_in_filename(tmp_path):
|
||||
gguf = tmp_path / "Qwen3.6-27B-MTP-Q4_K_M.gguf"
|
||||
gguf.write_bytes(b"")
|
||||
|
|
|
|||
|
|
@ -104,6 +104,9 @@ def _make_backend(effective_ctx = 98304, port = 51234):
|
|||
inst._port = port
|
||||
inst._effective_context_length = effective_ctx
|
||||
inst._context_length = 262144
|
||||
inst._effective_parallel_slots = 1
|
||||
inst._kv_cache_unified = False
|
||||
inst._kv_cache_context_total = None
|
||||
return inst
|
||||
|
||||
|
||||
|
|
@ -173,6 +176,31 @@ def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch):
|
|||
assert inst.context_length == 67584
|
||||
|
||||
|
||||
def test_props_keeps_total_cache_context_for_slot_preflight(monkeypatch):
|
||||
inst = _make_backend(effective_ctx = 32768)
|
||||
inst._effective_parallel_slots = 4
|
||||
_stub_props(
|
||||
monkeypatch,
|
||||
body = {"default_generation_settings": {"n_ctx": 8192}},
|
||||
)
|
||||
inst._reconcile_effective_ctx_with_server()
|
||||
assert inst._effective_context_length == 8192
|
||||
assert inst._kv_cache_context_total == 32768
|
||||
|
||||
|
||||
def test_props_does_not_multiply_unified_cache_context(monkeypatch):
|
||||
inst = _make_backend(effective_ctx = 32768)
|
||||
inst._effective_parallel_slots = 4
|
||||
inst._kv_cache_unified = True
|
||||
_stub_props(
|
||||
monkeypatch,
|
||||
body = {"default_generation_settings": {"n_ctx": 32768}},
|
||||
)
|
||||
inst._reconcile_effective_ctx_with_server()
|
||||
assert inst._effective_context_length == 32768
|
||||
assert inst._kv_cache_context_total == 32768
|
||||
|
||||
|
||||
def test_matching_ctx_is_left_alone(monkeypatch):
|
||||
inst = _make_backend(effective_ctx = 98304)
|
||||
_stub_props(
|
||||
|
|
|
|||
|
|
@ -221,6 +221,34 @@ def test_fingerprint_tracks_effective_context_length(tmp_path):
|
|||
assert backend._slot_launch_fingerprint() != before
|
||||
|
||||
|
||||
def test_fingerprint_tracks_swa_full_mode(tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
before = backend._slot_launch_fingerprint()
|
||||
backend._swa_full = True
|
||||
assert backend._slot_launch_fingerprint() != before
|
||||
|
||||
|
||||
def test_fingerprint_tracks_unified_cache_mode(tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
before = backend._slot_launch_fingerprint()
|
||||
backend._kv_cache_unified = True
|
||||
assert backend._slot_launch_fingerprint() != before
|
||||
|
||||
|
||||
def test_fingerprint_tracks_flash_attention_mode(tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
before = backend._slot_launch_fingerprint()
|
||||
backend._flash_attn_enabled = False
|
||||
assert backend._slot_launch_fingerprint() != before
|
||||
|
||||
|
||||
def test_fingerprint_tracks_effective_cache_types(tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
before = backend._slot_launch_fingerprint()
|
||||
backend._effective_cache_types = ("f32", "f16")
|
||||
assert backend._slot_launch_fingerprint() != before
|
||||
|
||||
|
||||
def test_gguf_file_identity_covers_split_shards(tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
first = tmp_path / "m-00001-of-00002.gguf"
|
||||
|
|
@ -444,6 +472,81 @@ def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path):
|
|||
assert backend.save_slots_for_resume() is None
|
||||
|
||||
|
||||
def test_save_estimate_uses_total_context_and_active_cache_settings(monkeypatch, tmp_path):
|
||||
backend = _resume_backend(tmp_path, n_slots = 4)
|
||||
backend._effective_context_length = 8192
|
||||
backend._kv_cache_context_total = 32768
|
||||
backend._sliding_window = 4096
|
||||
backend._swa_full = True
|
||||
backend._flash_attn_enabled = False
|
||||
backend._effective_cache_types = ("f32", "f16")
|
||||
calls = []
|
||||
|
||||
def estimate(ctx, cache_type, **kwargs):
|
||||
calls.append((ctx, cache_type, kwargs))
|
||||
return 0
|
||||
|
||||
backend._estimate_kv_cache_bytes = estimate
|
||||
_fake_disk(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
llama_cpp.httpx,
|
||||
"post",
|
||||
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
|
||||
raising = False,
|
||||
)
|
||||
|
||||
assert backend.save_slots_for_resume() is not None
|
||||
assert calls == [
|
||||
(
|
||||
32768,
|
||||
"f32",
|
||||
{
|
||||
"n_parallel": 4,
|
||||
"swa_full": True,
|
||||
"kv_unified": False,
|
||||
"n_ubatch": 512,
|
||||
"flash_attn": False,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_compact_swa_slot_save_is_skipped(monkeypatch, tmp_path):
|
||||
backend = _resume_backend(tmp_path)
|
||||
backend._sliding_window = 4096
|
||||
backend._kv_key_length = 256
|
||||
backend._kv_value_length = 256
|
||||
backend._swa_full = False
|
||||
backend._estimate_kv_cache_bytes = lambda *a, **k: (_ for _ in ()).throw(AssertionError)
|
||||
monkeypatch.setattr(
|
||||
llama_cpp.httpx,
|
||||
"post",
|
||||
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
|
||||
raising = False,
|
||||
)
|
||||
assert backend.save_slots_for_resume() is None
|
||||
|
||||
|
||||
def test_window_without_kv_dims_still_saves(monkeypatch, tmp_path):
|
||||
# phi3 reports a window but no key/value length, and llama.cpp runs it
|
||||
# non-SWA, so the compact-SWA skip must not catch it.
|
||||
backend = _resume_backend(tmp_path)
|
||||
backend._sliding_window = 262144
|
||||
backend._kv_key_length = None
|
||||
backend._kv_value_length = None
|
||||
backend._swa_full = False
|
||||
posted = []
|
||||
monkeypatch.setattr(
|
||||
llama_cpp.httpx,
|
||||
"post",
|
||||
lambda *a, **k: posted.append(a)
|
||||
or SimpleNamespace(status_code = 200, json = lambda: {"filename": "slot.bin"}),
|
||||
raising = False,
|
||||
)
|
||||
backend.save_slots_for_resume()
|
||||
assert posted
|
||||
|
||||
|
||||
def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path):
|
||||
# The GGUF/sidecars were swapped on disk after the server loaded them, so the
|
||||
# live KV belongs to the old weights: refuse to persist it (no POST at all).
|
||||
|
|
|
|||
|
|
@ -602,7 +602,7 @@ def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch):
|
|||
]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads)
|
||||
_patch_monotonic(monkeypatch, [200.0, 201.0, 203.0, 300.0, 400.0, 405.0, 405.0])
|
||||
_patch_monotonic(monkeypatch, [200.0, 201.0, 203.0, 300.0, 400.0, 405.0, 410.0])
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
return "Rendered HTML canvas: Done."
|
||||
|
|
@ -1495,6 +1495,7 @@ def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
|
|||
streams = [
|
||||
[_sse({"content": "I will use render_html now."}), _done()],
|
||||
[
|
||||
_sse({"reasoning_content": "I reconsidered the request."}),
|
||||
_sse({"content": "No tool is needed. Final answer: use a red square."}),
|
||||
_done(),
|
||||
],
|
||||
|
|
@ -1531,8 +1532,19 @@ def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
|
|||
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
|
||||
assert content_texts == [
|
||||
"I will use render_html now.",
|
||||
"No tool is needed. Final answer: use a red square.",
|
||||
(
|
||||
"<think>I reconsidered the request.</think>"
|
||||
"No tool is needed. Final answer: use a red square."
|
||||
),
|
||||
]
|
||||
summaries = [event for event in events if event.get("type") == "reasoning_summary"]
|
||||
assert len(summaries) == 1
|
||||
visible_answer_index = next(
|
||||
index
|
||||
for index, event in enumerate(events)
|
||||
if event.get("type") == "content" and "No tool is needed" in event.get("text", "")
|
||||
)
|
||||
assert visible_answer_index < events.index(summaries[0])
|
||||
assert len(payloads) == 2
|
||||
|
||||
|
||||
|
|
@ -1774,24 +1786,14 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
|
|||
streams = [
|
||||
[_sse({"content": "I will use render_html now."}), _done()],
|
||||
[
|
||||
_sse({"reasoning_content": "I should render the requested HTML."}),
|
||||
_sse(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_forced",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "render_html",
|
||||
"arguments": json.dumps(
|
||||
{
|
||||
"code": "<html><body>forced</body></html>",
|
||||
"title": "Forced",
|
||||
}
|
||||
),
|
||||
},
|
||||
}
|
||||
]
|
||||
"content": (
|
||||
'<tool_call>{"name":"render_html","arguments":'
|
||||
'{"code":"<html><body>forced</body></html>",'
|
||||
'"title":"Forced"}}</tool_call>'
|
||||
)
|
||||
}
|
||||
),
|
||||
_done(),
|
||||
|
|
@ -1835,6 +1837,7 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
|
|||
assert len(calls) == 1
|
||||
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
|
||||
assert content_texts == ["I will use render_html now.", "Final note after tool."]
|
||||
assert not any(event.get("type") == "reasoning_summary" for event in events)
|
||||
assert len(payloads) == 3
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -113,6 +113,11 @@ def test_value_with_equals_form_passes_through():
|
|||
assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"]
|
||||
|
||||
|
||||
def test_managed_long_flag_underscore_alias_is_rejected():
|
||||
with pytest.raises(ValueError, match = "slot-save-path"):
|
||||
validate_extra_args(["--slot_save_path", "/tmp/slots"])
|
||||
|
||||
|
||||
def test_non_flag_token_passes_through():
|
||||
# Bare positionals are passed through; llama-server can reject them.
|
||||
assert validate_extra_args(["foo"]) == ["foo"]
|
||||
|
|
|
|||
|
|
@ -76,7 +76,9 @@ from core.inference.llama_cpp import ( # noqa: E402
|
|||
_extra_args_spec_draft_n_max,
|
||||
_effective_tensor_parallel,
|
||||
_env_main_cache_type_for_budget,
|
||||
_effective_main_cache_types,
|
||||
_extra_args_main_cache_type_for_budget,
|
||||
_flash_attn_enabled_from_args,
|
||||
_kv_bytes_per_elem,
|
||||
_tensor_parallel_matches_loaded,
|
||||
)
|
||||
|
|
@ -132,6 +134,7 @@ class _StubDrafter:
|
|||
|
||||
def __init__(self, kv_per_token):
|
||||
self._kv_per_token = kv_per_token
|
||||
self._architecture = "gemma3"
|
||||
|
||||
def _can_estimate_kv(self):
|
||||
return True
|
||||
|
|
@ -177,6 +180,14 @@ class TestEmbeddedDraftKv:
|
|||
two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536)
|
||||
assert two == pytest.approx(2 * one)
|
||||
|
||||
def test_unaligned_context_follows_runtime_stream_padding(self):
|
||||
b = _make_backend()
|
||||
bytes_per_cell = b._mtp_draft_kv_bytes(256) // 256
|
||||
unified = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = True)
|
||||
separate = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = False)
|
||||
assert unified == 5120 * bytes_per_cell
|
||||
assert separate == 5376 * bytes_per_cell
|
||||
|
||||
def test_embedded_draft_kv_floored_at_f16(self):
|
||||
# The embedded MTP head is one layer, so llama.cpp's quantized-KV
|
||||
# overhead is not amortized: a quantized draft KV fits LESS context than
|
||||
|
|
@ -201,6 +212,15 @@ class TestEmbeddedDraftKv:
|
|||
both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16")
|
||||
assert both_q4 == k_only == both_f16 # floored at f16, never under-reserved
|
||||
|
||||
def test_flash_attn_off_uses_model_wide_v_width(self):
|
||||
b = _make_backend(n_layers = 2)
|
||||
b._n_kv_heads_by_layer = [4, 1]
|
||||
b._sliding_window_pattern = [False, True]
|
||||
b._kv_value_length_swa = 2048
|
||||
ctx = 4096
|
||||
expected_per_cell = 4 * 256 * 2 + 1 * 2048 * 2
|
||||
assert b._mtp_draft_kv_bytes(ctx, flash_attn = False) == ctx * expected_per_cell
|
||||
|
||||
def test_none_when_dims_missing(self):
|
||||
assert _make_backend(nextn = 0)._mtp_draft_kv_bytes(65536) is None
|
||||
assert _make_backend(kv_key_length = None)._mtp_draft_kv_bytes(65536) is None
|
||||
|
|
@ -232,6 +252,30 @@ class TestSeparateDrafter:
|
|||
c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf")
|
||||
assert c == pytest.approx(4 * a)
|
||||
|
||||
def test_gemma4_assistant_shares_target_kv(self, monkeypatch):
|
||||
b = _make_backend(nextn = None)
|
||||
stub = _StubDrafter(kv_per_token = 2000)
|
||||
stub._architecture = "gemma4-assistant"
|
||||
monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub)
|
||||
|
||||
assert (
|
||||
b._mtp_draft_kv_bytes(
|
||||
65536,
|
||||
drafter_path = "/m/mtp-gemma4.gguf",
|
||||
swa_full = True,
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert (
|
||||
b._estimate_mtp_overhead_bytes(
|
||||
65536,
|
||||
drafter_path = "/m/mtp-gemma4.gguf",
|
||||
draft_weights_bytes = GIB,
|
||||
swa_full = True,
|
||||
)
|
||||
== GIB
|
||||
)
|
||||
|
||||
def test_drafter_kv_scales_with_parallel_slots(self, monkeypatch):
|
||||
# The drafter is served under the same --parallel slots as the main model,
|
||||
# so a sliding-window drafter's KV grows per slot; the reserve must thread
|
||||
|
|
@ -398,6 +442,7 @@ class TestExtraArgsMtpDetection:
|
|||
(["--spec-type", "mtp"], True),
|
||||
(["--spec-type", "ngram-mod,draft-mtp"], True),
|
||||
(["--spec-type=draft-mtp"], True),
|
||||
(["--spec_type=draft-mtp"], True),
|
||||
(["--spec-type", "ngram-mod"], False),
|
||||
(["--spec-default"], False),
|
||||
(["-c", "131072"], False),
|
||||
|
|
@ -579,6 +624,7 @@ class TestExtraArgsMtpDetection:
|
|||
(["--spec-draft-ngl", "0"], True),
|
||||
(["-ngld", "0"], True),
|
||||
(["--spec-draft-ngl=0"], True),
|
||||
(["--spec_draft_ngl=0"], True),
|
||||
(["--n-gpu-layers-draft", "0"], True),
|
||||
(["--spec-draft-ngl", "20"], False),
|
||||
(["--spec-draft-device", "none"], True),
|
||||
|
|
@ -623,6 +669,7 @@ class TestExtraArgsMtpDetection:
|
|||
[
|
||||
(["--spec-draft-n-max", "4"], 4),
|
||||
(["--spec-draft-n-max=6"], 6),
|
||||
(["--spec_draft_n_max=6"], 6),
|
||||
(["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"], 3),
|
||||
(["--spec-draft-n-max", "2", "--spec-draft-n-max", "5"], 5), # last wins
|
||||
(["--spec-draft-n-max", "notanint"], None),
|
||||
|
|
@ -644,6 +691,7 @@ class TestExtraArgsMtpDetection:
|
|||
(["--spec-draft-model", "/m/draft.gguf"], "/m/draft.gguf"),
|
||||
(["-md", "/m/draft.gguf"], "/m/draft.gguf"),
|
||||
(["--model-draft=/m/draft.gguf"], "/m/draft.gguf"),
|
||||
(["--model_draft=/m/draft.gguf"], "/m/draft.gguf"),
|
||||
(["--model-draft", "--spec-type"], None),
|
||||
(["-c", "4096"], None),
|
||||
(None, None),
|
||||
|
|
@ -689,6 +737,7 @@ class TestExtraArgsMtpDetection:
|
|||
(["--cache-type-v-draft", "q4_0"], (None, "q4_0")), # K stays f16, V only
|
||||
(["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")),
|
||||
(["--cache-type-k-draft=q8_0"], ("q8_0", None)),
|
||||
(["--cache_type_k_draft=q8_0"], ("q8_0", None)),
|
||||
(["--cache-type-k", "q8_0"], (None, None)), # main type, not draft
|
||||
(["-c", "4096"], (None, None)),
|
||||
(None, (None, None)),
|
||||
|
|
@ -717,8 +766,17 @@ class TestExtraArgsMtpDetection:
|
|||
"args,expected",
|
||||
[
|
||||
(["--ubatch-size", "1024"], 1024),
|
||||
(["-ub", "4096"], 4096),
|
||||
(["-ub", "4096"], 2048),
|
||||
(["--ubatch-size", "0"], 2048),
|
||||
(["--batch-size", "256", "--ubatch-size", "0"], 256),
|
||||
(["--batch-size", "-1"], 512),
|
||||
(["--ubatch-size", "-1"], 2048),
|
||||
(["--ubatch-size=512"], 512),
|
||||
(["--ubatch_size=512"], 512),
|
||||
(["--batch-size", "256"], 256),
|
||||
(["--batch_size=256"], 256),
|
||||
(["-b", "256", "-ub", "1024"], 256),
|
||||
(["-b", "4096"], 512),
|
||||
(["--ubatch", "2048"], None), # not a real llama-server flag; ignore it
|
||||
(["-c", "4096"], None),
|
||||
(None, None),
|
||||
|
|
@ -727,12 +785,76 @@ class TestExtraArgsMtpDetection:
|
|||
def test_n_ubatch(self, args, expected):
|
||||
assert _extra_args_n_ubatch(args, env = {}) == expected
|
||||
|
||||
def test_n_ubatch_signed_values_cap_at_context(self):
|
||||
assert (
|
||||
_extra_args_n_ubatch(
|
||||
["--batch-size", "-1", "--ubatch-size", "-1"],
|
||||
env = {},
|
||||
n_ctx = 4096,
|
||||
)
|
||||
== 4096
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args,expected",
|
||||
[
|
||||
(None, True),
|
||||
(["--flash-attn", "off"], False),
|
||||
(["--flash-attn", "disabled"], False),
|
||||
(["--flash-attn", "false"], False),
|
||||
(["--flash-attn", "0"], False),
|
||||
(["--flash-attn=off"], False),
|
||||
(["--flash-attn=disabled"], False),
|
||||
(["--flash-attn=false"], False),
|
||||
(["--flash-attn=0"], False),
|
||||
(["--flash_attn", "off"], False),
|
||||
(["-fa", "off", "--flash-attn", "auto"], True),
|
||||
(["-fa", "off", "--flash-attn", "-1"], True),
|
||||
(["-fa", "off", "--flash-attn", "enabled"], True),
|
||||
(["-fa", "off", "--flash-attn=true"], True),
|
||||
(["-fa", "off", "--flash-attn=1"], True),
|
||||
(["--flash-attn", "off", "-fa"], True),
|
||||
],
|
||||
)
|
||||
def test_flash_attn_last_value_wins(self, args, expected):
|
||||
assert _flash_attn_enabled_from_args(args) is expected
|
||||
|
||||
def test_effective_main_cache_types_follow_env_then_cli(self):
|
||||
env = {
|
||||
"LLAMA_ARG_CACHE_TYPE_K": "f32",
|
||||
"LLAMA_ARG_CACHE_TYPE_V": "q4_0",
|
||||
}
|
||||
assert _effective_main_cache_types([], env) == ("f32", "q4_0")
|
||||
assert _effective_main_cache_types(["--cache-type-v", "f16"], env) == ("f32", "f16")
|
||||
|
||||
def test_n_ubatch_env_fallback(self):
|
||||
# The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve.
|
||||
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096
|
||||
# Environment values apply first, then each command-line option overrides
|
||||
# its own axis before llama.cpp caps ubatch at batch size.
|
||||
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 2048
|
||||
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_BATCH": "256"}) == 256
|
||||
assert (
|
||||
_extra_args_n_ubatch(
|
||||
[],
|
||||
env = {
|
||||
"LLAMA_ARG_BATCH": "1024",
|
||||
"LLAMA_ARG_UBATCH": "4096",
|
||||
},
|
||||
)
|
||||
== 1024
|
||||
)
|
||||
assert (
|
||||
_extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024
|
||||
) # CLI wins
|
||||
assert (
|
||||
_extra_args_n_ubatch(
|
||||
["-b", "1024"],
|
||||
env = {
|
||||
"LLAMA_ARG_BATCH": "256",
|
||||
"LLAMA_ARG_UBATCH": "4096",
|
||||
},
|
||||
)
|
||||
== 1024
|
||||
)
|
||||
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None
|
||||
|
||||
def test_env_main_cache_type_for_budget(self):
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ def _backend(
|
|||
vocab = 248320,
|
||||
embd = 5120,
|
||||
kv_fixed_mib = 0,
|
||||
kv_calls = None,
|
||||
):
|
||||
"""Backend with the dims the compute buffer reads; KV mocked to a fixed size so the
|
||||
only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15)."""
|
||||
|
|
@ -43,7 +44,17 @@ def _backend(
|
|||
b._vocab_size = vocab
|
||||
b._embedding_length = embd
|
||||
b._key_length_mla = None
|
||||
b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB
|
||||
|
||||
def estimate(
|
||||
ctx,
|
||||
t = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kv_calls is not None:
|
||||
kv_calls.append(kwargs)
|
||||
return kv_fixed_mib * MIB
|
||||
|
||||
b._estimate_kv_cache_bytes = estimate
|
||||
b._can_estimate_kv = lambda: True
|
||||
return b
|
||||
|
||||
|
|
@ -55,6 +66,7 @@ def _run(
|
|||
gpus,
|
||||
total_by_idx,
|
||||
overhead_mib = 0,
|
||||
swa_full = False,
|
||||
):
|
||||
return b._slots_that_fit_on_gpu(
|
||||
n_parallel,
|
||||
|
|
@ -66,7 +78,8 @@ def _run(
|
|||
FRAC,
|
||||
int(overhead_mib * MIB),
|
||||
1,
|
||||
512,
|
||||
n_ubatch = 512,
|
||||
swa_full = swa_full,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -113,3 +126,16 @@ class TestSlotsThatFitOnGpu:
|
|||
# base 19500 (= 22500 total at par-independent terms) the same par3 fit holds.
|
||||
gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576})
|
||||
assert use_fit is False and slots == 3
|
||||
|
||||
def test_swa_full_is_used_for_every_candidate(self):
|
||||
calls = []
|
||||
_run(
|
||||
_backend(kv_calls = calls),
|
||||
4,
|
||||
22500,
|
||||
[(0, 24576)],
|
||||
{0: 24576},
|
||||
swa_full = True,
|
||||
)
|
||||
assert calls
|
||||
assert all(call["swa_full"] is True for call in calls)
|
||||
|
|
|
|||
|
|
@ -209,6 +209,13 @@ def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, reque
|
|||
assert _target_state(_loaded_backend(loaded), requested) is False
|
||||
|
||||
|
||||
def test_already_in_target_state_reloads_when_swa_full_env_changes(monkeypatch):
|
||||
backend = _loaded_backend(False)
|
||||
backend._swa_full = False
|
||||
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
|
||||
assert _target_state(backend, False) is False
|
||||
|
||||
|
||||
def test_already_in_target_state_reconciles_split_mode_extras():
|
||||
# Tensor engaged via --split-mode in extras (boolean omitted/default False)
|
||||
# must match a server already running tensor mode -- no spurious reload.
|
||||
|
|
|
|||
|
|
@ -663,6 +663,29 @@ def test_tensor_off_echo_preserves_multi_gpu_fallback():
|
|||
)
|
||||
|
||||
|
||||
def test_route_dedupe_reloads_when_swa_full_env_changes(monkeypatch):
|
||||
from models.inference import LoadRequest
|
||||
|
||||
inference_routes = _load_inference_routes_module()
|
||||
backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False)
|
||||
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
|
||||
|
||||
request = LoadRequest(model_path = "owner/repo")
|
||||
assert inference_routes._request_matches_loaded_settings(request, backend) is False
|
||||
|
||||
|
||||
def test_route_dedupe_ignores_swa_full_for_diffusion(monkeypatch):
|
||||
from models.inference import LoadRequest
|
||||
|
||||
inference_routes = _load_inference_routes_module()
|
||||
backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False)
|
||||
backend._is_diffusion = True
|
||||
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
|
||||
|
||||
request = LoadRequest(model_path = "owner/repo")
|
||||
assert inference_routes._request_matches_loaded_settings(request, backend) is True
|
||||
|
||||
|
||||
def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback():
|
||||
"""Tensor intent can be dropped via extras too: an explicit --split-mode layer
|
||||
matches the stored fallback extras but must still reload (reviewer.py P1, #6659)."""
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@
|
|||
"build": "tsc -b && vite build",
|
||||
"lint": "eslint .",
|
||||
"preview": "vite preview",
|
||||
"typecheck": "tsc -b --pretty false",
|
||||
"test": "node --experimental-strip-types --test \"tests/**/*.test.ts\"",
|
||||
"typecheck": "tsc -b --pretty false && tsc -p tsconfig.test.json --pretty false",
|
||||
"i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts",
|
||||
"biome:check": "biome check",
|
||||
"biome:fix": "biome check --write"
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -11,6 +11,7 @@ import {
|
|||
CollapsibleContent,
|
||||
CollapsibleTrigger,
|
||||
} from "@/components/ui/collapsible";
|
||||
import { resolveReasoningGroupDuration } from "@/features/chat";
|
||||
import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
|
|
@ -339,9 +340,11 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
|
|||
});
|
||||
|
||||
const persistedDuration = useAuiState(({ message }) => {
|
||||
const d = (message.metadata?.custom as Record<string, unknown>)
|
||||
?.reasoningDuration;
|
||||
return typeof d === "number" ? d : 0;
|
||||
return resolveReasoningGroupDuration(
|
||||
message.parts,
|
||||
startIndex,
|
||||
message.metadata?.custom as Record<string, unknown> | undefined,
|
||||
);
|
||||
});
|
||||
|
||||
const [manualOpen, setManualOpen] = useState(false);
|
||||
|
|
@ -412,7 +415,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
|
|||
className="min-w-0 flex-1"
|
||||
active={isReasoningStreaming}
|
||||
// Prefer server timing when available.
|
||||
duration={persistedDuration || duration}
|
||||
duration={persistedDuration ?? duration}
|
||||
/>
|
||||
<div className="flex w-16 shrink-0 justify-end">
|
||||
{isOpen && !isReasoningStreaming && (
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
|
|||
import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
|
||||
import { ChatDictationBar } from "@/components/assistant-ui/chat-dictation-bar";
|
||||
import {
|
||||
pasteClipboardFiles,
|
||||
isStudioDictationAvailable,
|
||||
notifyStudioDictationUnavailable,
|
||||
} from "@/features/chat";
|
||||
|
|
@ -177,6 +178,7 @@ import {
|
|||
type ChangeEvent,
|
||||
type ComponentProps,
|
||||
type CompositionEvent,
|
||||
type ClipboardEvent,
|
||||
type FC,
|
||||
type KeyboardEvent,
|
||||
type DragEvent as ReactDragEvent,
|
||||
|
|
@ -1528,6 +1530,24 @@ const Composer: FC<{
|
|||
);
|
||||
const { inputProps, isComposing, isComposingRef } =
|
||||
useImeComposerInputHandlers({ submitOnEnter: true });
|
||||
const handleFilePaste = useCallback(
|
||||
(event: ClipboardEvent<HTMLTextAreaElement>) => {
|
||||
pasteClipboardFiles(
|
||||
event,
|
||||
async (files) => {
|
||||
await Promise.all(
|
||||
files.map((file) => aui.composer().addAttachment(file)),
|
||||
);
|
||||
},
|
||||
() =>
|
||||
toast.error("Could not paste files.", {
|
||||
description: "The clipboard item is unsupported, unreadable, or over 20 MB.",
|
||||
}),
|
||||
);
|
||||
},
|
||||
[aui],
|
||||
);
|
||||
|
||||
const composerText = useAuiState(({ composer }) => composer.text);
|
||||
// Expand only once the input wraps to a second line, not on first keystroke.
|
||||
// Latch until cleared so it can't flip-flop at the wrap boundary.
|
||||
|
|
@ -2021,6 +2041,8 @@ const Composer: FC<{
|
|||
// no effect on Latin / CJK / Devanagari.
|
||||
dir="auto"
|
||||
{...inputProps}
|
||||
addAttachmentOnPaste={false}
|
||||
onPaste={handleFilePaste}
|
||||
/>
|
||||
<ComposerRightControls
|
||||
disabled={
|
||||
|
|
|
|||
|
|
@ -86,9 +86,15 @@ import {
|
|||
} from "../utils/last-local-model-load";
|
||||
import { getImageInputUnavailableReason } from "../utils/image-input-support";
|
||||
import {
|
||||
hasClosedThinkTag,
|
||||
extractDeltaText,
|
||||
hasUnclosedThinkTag,
|
||||
parseAssistantContent,
|
||||
} from "../utils/parse-assistant-content";
|
||||
import {
|
||||
countReasoningGroups,
|
||||
createReasoningDurationTracker,
|
||||
lastReasoningGroupTextLength,
|
||||
} from "../utils/reasoning-duration";
|
||||
import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
|
||||
import {
|
||||
generateAudio,
|
||||
|
|
@ -617,67 +623,6 @@ function estimateTokenCount(text: string): number | undefined {
|
|||
return Math.max(1, Math.round(trimmed.length / 4));
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize a streamed `delta.content` to a plain text string.
|
||||
*
|
||||
* OpenAI Chat Completions originally typed `delta.content` as a string, but
|
||||
* some providers now emit an array of structured content parts; concatenating
|
||||
* those directly would stringify each as `[object Object]`. This guards that.
|
||||
*
|
||||
* Handled part shapes:
|
||||
* { type: "text" | "output_text", text | content: "..." } → text body
|
||||
* { type: "thinking" | "reasoning", thinking | text: "..." } → wrapped as
|
||||
* inline `<think>...</think>` so `parseAssistantContent` lifts it into
|
||||
* a reasoning part (else Mistral magistral and similar reasoning-part
|
||||
* providers lose their thinking panel).
|
||||
*
|
||||
* Unknown part types are skipped — better to drop a stray field than
|
||||
* stringify an object into the rendered chat.
|
||||
*/
|
||||
function extractDeltaText(delta: unknown): string {
|
||||
const extractReasoningText = (payload: unknown): string => {
|
||||
if (typeof payload === "string") return payload;
|
||||
if (Array.isArray(payload)) {
|
||||
return payload.map((item) => extractReasoningText(item)).join("");
|
||||
}
|
||||
if (!payload || typeof payload !== "object") return "";
|
||||
|
||||
const obj = payload as Record<string, unknown>;
|
||||
for (const key of ["thinking", "text", "content", "reasoning", "summary"]) {
|
||||
if (key in obj) {
|
||||
const text = extractReasoningText(obj[key]);
|
||||
if (text) return text;
|
||||
}
|
||||
}
|
||||
return "";
|
||||
};
|
||||
|
||||
if (typeof delta === "string") return delta;
|
||||
if (!Array.isArray(delta)) return "";
|
||||
let out = "";
|
||||
for (const part of delta) {
|
||||
if (typeof part === "string") {
|
||||
out += part;
|
||||
continue;
|
||||
}
|
||||
if (!part || typeof part !== "object") continue;
|
||||
const obj = part as {
|
||||
type?: string;
|
||||
text?: string;
|
||||
content?: string;
|
||||
thinking?: string;
|
||||
};
|
||||
if (obj.type === "text" || obj.type === "output_text") {
|
||||
if (typeof obj.text === "string") out += obj.text;
|
||||
else if (typeof obj.content === "string") out += obj.content;
|
||||
} else if (obj.type === "thinking" || obj.type === "reasoning") {
|
||||
const thinking = extractReasoningText(obj);
|
||||
if (thinking) out += `<think>${thinking}</think>`;
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function buildTiming(
|
||||
streamStartTime: number,
|
||||
totalChunks: number,
|
||||
|
|
@ -1562,6 +1507,8 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
// The safetensors fallback omits both fields and uses HF auto-placement.
|
||||
gpu_ids?: number[];
|
||||
gpu_memory_mode?: "auto" | "manual";
|
||||
cache_type_kv?: string | null;
|
||||
tensor_parallel?: boolean | null;
|
||||
}): Promise<boolean> {
|
||||
const validation = await validateModel({
|
||||
...payload,
|
||||
|
|
@ -1650,6 +1597,8 @@ async function autoLoadSmallestModel(): Promise<{
|
|||
max_seq_length: fitMaxSeqLength,
|
||||
is_lora: false,
|
||||
gguf_variant: candidate.ggufVariant,
|
||||
cache_type_kv: config.kvCacheDtype,
|
||||
tensor_parallel: config.tensorParallel,
|
||||
// The same remembered-derived GPU pick the load below sends.
|
||||
...(candidate.kind === "gguf"
|
||||
? {
|
||||
|
|
@ -2939,8 +2888,7 @@ export function createOpenAIStreamAdapter(
|
|||
owner: serverCancel,
|
||||
});
|
||||
let cumulativeText = "";
|
||||
let reasoningStartAt: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
const reasoningDurationTracker = createReasoningDurationTracker();
|
||||
// True while wrapping a `delta.reasoning_content` stream in
|
||||
// <think>...</think> for parseAssistantContent. Lives outside the
|
||||
// SSE loop because the close tag fires when content arrives.
|
||||
|
|
@ -3080,9 +3028,11 @@ export function createOpenAIStreamAdapter(
|
|||
return merged;
|
||||
};
|
||||
const closeReasoningContent = () => {
|
||||
if (!reasoningContentOpen) return;
|
||||
cumulativeText += "</think>";
|
||||
reasoningContentOpen = false;
|
||||
if (reasoningContentOpen) {
|
||||
cumulativeText += "</think>";
|
||||
reasoningContentOpen = false;
|
||||
}
|
||||
reasoningDurationTracker.finishGroup();
|
||||
};
|
||||
// Anthropic document_citations payload, converted to Sources-panel
|
||||
// parts at end-of-stream so inline [N] markers have matching entries.
|
||||
|
|
@ -3638,8 +3588,9 @@ export function createOpenAIStreamAdapter(
|
|||
const reasoningMs = (
|
||||
chunk as { _reasoningDurationMs?: number } | null | undefined
|
||||
)?._reasoningDurationMs;
|
||||
if (typeof reasoningMs === "number" && Number.isFinite(reasoningMs)) {
|
||||
reasoningDuration = Math.max(0, Math.round(reasoningMs / 1000));
|
||||
if (
|
||||
reasoningDurationTracker.recordServerDuration(reasoningMs)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -3783,7 +3734,7 @@ export function createOpenAIStreamAdapter(
|
|||
totalChunks,
|
||||
firstTokenTime,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
custom: reasoningDurationTracker.metadata(),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -4078,7 +4029,7 @@ export function createOpenAIStreamAdapter(
|
|||
totalChunks,
|
||||
firstTokenTime,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
custom: reasoningDurationTracker.metadata(),
|
||||
},
|
||||
};
|
||||
continue;
|
||||
|
|
@ -4117,7 +4068,10 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
const rawDelta = chunk.choices?.[0]?.delta?.content;
|
||||
// Normalize structured delta.content (mistral magistral).
|
||||
const delta = extractDeltaText(rawDelta);
|
||||
const {
|
||||
text: delta,
|
||||
structuredReasoningContinues,
|
||||
} = extractDeltaText(rawDelta);
|
||||
// Latest Gemini text-part thoughtSignature for next-turn replay.
|
||||
const deltaExtraContent = (
|
||||
chunk.choices?.[0]?.delta as
|
||||
|
|
@ -4271,7 +4225,7 @@ export function createOpenAIStreamAdapter(
|
|||
totalChunks,
|
||||
firstTokenTime,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
custom: reasoningDurationTracker.metadata(),
|
||||
},
|
||||
};
|
||||
continue;
|
||||
|
|
@ -4288,6 +4242,7 @@ export function createOpenAIStreamAdapter(
|
|||
|
||||
if (reasoning) {
|
||||
if (!reasoningContentOpen) {
|
||||
reasoningDurationTracker.startGroup();
|
||||
cumulativeText += `<think>${reasoning}`;
|
||||
reasoningContentOpen = true;
|
||||
} else {
|
||||
|
|
@ -4295,7 +4250,9 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
}
|
||||
if (delta) {
|
||||
closeReasoningContent();
|
||||
if (reasoningContentOpen) {
|
||||
closeReasoningContent();
|
||||
}
|
||||
cumulativeText += delta;
|
||||
}
|
||||
// Strip a trailing ${...} template-literal fragment from
|
||||
|
|
@ -4306,35 +4263,48 @@ export function createOpenAIStreamAdapter(
|
|||
"",
|
||||
);
|
||||
}
|
||||
const textParts = parseAssistantContent(cumulativeText);
|
||||
const assistantContent = buildAssistantContent(cumulativeText);
|
||||
|
||||
// Fallback when no server-side reasoning_summary arrives.
|
||||
const parsedReasoningGroupCount =
|
||||
countReasoningGroups(assistantContent);
|
||||
if (
|
||||
textParts.some((part) => part.type === "reasoning") &&
|
||||
!reasoningStartAt
|
||||
parsedReasoningGroupCount >
|
||||
reasoningDurationTracker.groupCount
|
||||
) {
|
||||
reasoningStartAt = Date.now();
|
||||
}
|
||||
if (
|
||||
hasClosedThinkTag(cumulativeText) &&
|
||||
reasoningStartAt &&
|
||||
!reasoningDuration
|
||||
) {
|
||||
reasoningDuration = Math.round(
|
||||
(Date.now() - reasoningStartAt) / 1000,
|
||||
reasoningDurationTracker.startGroup(
|
||||
parsedReasoningGroupCount - 1,
|
||||
);
|
||||
}
|
||||
if (parsedReasoningGroupCount > 0) {
|
||||
// Providers that close every reasoning block atomically
|
||||
// (structured parts wrapped as <think>..</think>) end the group
|
||||
// on each chunk. Reopen while the reasoning text is still
|
||||
// growing so the timer spans the whole pass.
|
||||
reasoningDurationTracker.resumeGroup(
|
||||
parsedReasoningGroupCount - 1,
|
||||
lastReasoningGroupTextLength(assistantContent),
|
||||
);
|
||||
}
|
||||
if (
|
||||
reasoningDurationTracker.hasActiveGroup &&
|
||||
!reasoningContentOpen &&
|
||||
!structuredReasoningContinues &&
|
||||
!hasUnclosedThinkTag(cumulativeText)
|
||||
) {
|
||||
reasoningDurationTracker.finishGroup();
|
||||
}
|
||||
|
||||
if (textParts.length > 0 || toolCallParts.length > 0) {
|
||||
if (assistantContent.length > 0) {
|
||||
yield {
|
||||
content: buildAssistantContent(cumulativeText),
|
||||
content: assistantContent,
|
||||
metadata: {
|
||||
timing: buildTiming(
|
||||
streamStartTime,
|
||||
totalChunks,
|
||||
firstTokenTime,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
custom: reasoningDurationTracker.metadata(),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -4437,12 +4407,7 @@ export function createOpenAIStreamAdapter(
|
|||
);
|
||||
|
||||
// Finalize reasoning-only streams.
|
||||
if (reasoningStartAt && !reasoningDuration) {
|
||||
reasoningDuration = Math.max(
|
||||
0,
|
||||
Math.round((Date.now() - reasoningStartAt) / 1000),
|
||||
);
|
||||
}
|
||||
reasoningDurationTracker.finishGroup();
|
||||
yield {
|
||||
content: [
|
||||
...buildAssistantContent(cumulativeText),
|
||||
|
|
@ -4452,7 +4417,7 @@ export function createOpenAIStreamAdapter(
|
|||
metadata: {
|
||||
timing: finalTiming,
|
||||
custom: {
|
||||
reasoningDuration,
|
||||
...reasoningDurationTracker.metadata(),
|
||||
// Persisted refusal flag driving the two-pass prune.
|
||||
anthropicRefusal: anthropicRefusalSeen || undefined,
|
||||
serverTimings: meta?.timings ?? undefined,
|
||||
|
|
@ -4511,6 +4476,30 @@ export function createOpenAIStreamAdapter(
|
|||
});
|
||||
}
|
||||
}
|
||||
if (!abortSignal.aborted) {
|
||||
closeReasoningContent();
|
||||
const partialContent = buildAssistantContent(cumulativeText);
|
||||
if (partialContent.length > 0) {
|
||||
const partialTiming = buildTiming(
|
||||
streamStartTime,
|
||||
totalChunks,
|
||||
firstTokenTime,
|
||||
Date.now() - streamStartTime,
|
||||
estimateTokenCount(cumulativeText),
|
||||
toolCallParts.length,
|
||||
);
|
||||
yield {
|
||||
content: partialContent,
|
||||
metadata: {
|
||||
timing: partialTiming,
|
||||
custom: {
|
||||
...reasoningDurationTracker.metadata(),
|
||||
timing: partialTiming,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
throw err;
|
||||
} finally {
|
||||
runSignal.removeEventListener("abort", onAbortCancel);
|
||||
|
|
|
|||
|
|
@ -185,6 +185,8 @@ export async function validateModel(
|
|||
// /load. Default placement is sized against the selected GPUs.
|
||||
max_seq_length: payload.max_seq_length,
|
||||
load_in_4bit: payload.load_in_4bit,
|
||||
cache_type_kv: payload.cache_type_kv ?? null,
|
||||
tensor_parallel: payload.tensor_parallel ?? false,
|
||||
gpu_ids: payload.gpu_ids,
|
||||
// Manual placement is an explicit override: Auto layers use llama.cpp
|
||||
// --fit, while a pinned layer count is owned by the user. Tell validate
|
||||
|
|
|
|||
|
|
@ -824,6 +824,8 @@ export function useChatModelRuntime() {
|
|||
load_in_4bit: true,
|
||||
is_lora: isLora,
|
||||
gguf_variant: ggufVariant ?? null,
|
||||
cache_type_kv: loadKvCacheDtype,
|
||||
tensor_parallel: loadTensorParallel,
|
||||
gpu_ids: validateGpuIds ?? undefined,
|
||||
...(isGguf
|
||||
? {
|
||||
|
|
|
|||
|
|
@ -88,8 +88,10 @@ export { StopRunningChatsDialog } from "./components/stop-running-chats-dialog";
|
|||
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
|
||||
export type { ProjectRecord } from "./types";
|
||||
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
|
||||
export { pasteClipboardFiles } from "./utils/clipboard-files";
|
||||
export { listStoredChatThreads } from "./utils/chat-history-storage";
|
||||
export { emitChatAttachmentDeleted } from "./utils/chat-attachment-events";
|
||||
export { resolveReasoningGroupDuration } from "./utils/reasoning-duration";
|
||||
export { ArtifactCard } from "./artifacts/artifact-card";
|
||||
export { ResearchMessage } from "./components/research-message";
|
||||
export {
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ import { isTauri } from "@/lib/api-base";
|
|||
import { isDownloadCancelled } from "@/lib/native-files";
|
||||
import { isMultimodalResponse } from "./types/api";
|
||||
import { getImageInputUnavailableReason } from "./utils/image-input-support";
|
||||
import { pasteClipboardFiles } from "./utils/clipboard-files";
|
||||
import { useAui } from "@assistant-ui/react";
|
||||
import {
|
||||
ArrowUpIcon,
|
||||
|
|
@ -118,6 +119,7 @@ import {
|
|||
} from "./provider-capabilities";
|
||||
import {
|
||||
type CompositionEvent,
|
||||
type ClipboardEvent,
|
||||
type FC,
|
||||
type KeyboardEvent,
|
||||
type MutableRefObject,
|
||||
|
|
@ -834,7 +836,7 @@ export function SharedComposer({
|
|||
}, [text]);
|
||||
|
||||
const addFiles = useCallback(
|
||||
(files: FileList | null) => {
|
||||
(files: FileList | readonly File[] | null) => {
|
||||
if (!files?.length) return;
|
||||
const next: PendingImage[] = [];
|
||||
let droppedImageForUnavailable = false;
|
||||
|
|
@ -866,6 +868,29 @@ export function SharedComposer({
|
|||
[setPendingAudioStore, attachUnavailableReason],
|
||||
);
|
||||
|
||||
const handleFilePaste = useCallback(
|
||||
(event: ClipboardEvent<HTMLTextAreaElement>) => {
|
||||
pasteClipboardFiles(
|
||||
event,
|
||||
async (files) => {
|
||||
const supported = files.some(
|
||||
(file) =>
|
||||
(file.type.match(/^audio\//i) && file.size <= MAX_AUDIO_SIZE) ||
|
||||
(file.type.match(/^image\/(jpeg|png|webp|gif)$/i) &&
|
||||
file.size <= MAX_IMAGE_SIZE),
|
||||
);
|
||||
if (!supported) throw new Error("Unsupported compare attachment");
|
||||
addFiles(files);
|
||||
},
|
||||
() =>
|
||||
toast.error("Could not paste files.", {
|
||||
description: "Compare supports images and audio within the attachment size limits.",
|
||||
}),
|
||||
);
|
||||
},
|
||||
[addFiles],
|
||||
);
|
||||
|
||||
const removePendingImage = useCallback((id: string) => {
|
||||
setPendingImages((prev) => prev.filter((p) => p.id !== id));
|
||||
}, []);
|
||||
|
|
@ -1097,6 +1122,8 @@ export function SharedComposer({
|
|||
gguf_variant: sel.ggufVariant ?? null,
|
||||
trust_remote_code: loadTrustRemoteCode,
|
||||
chat_template_override: effectiveChatTemplateOverride,
|
||||
cache_type_kv: ownConfig.kvCacheDtype ?? null,
|
||||
tensor_parallel: effectiveTensorParallel,
|
||||
// Scope the validate to the picked GPUs. GGUF-only, like the load
|
||||
// below: a non-GGUF target must not inherit a hidden GGUF GPU pick.
|
||||
...(targetIsGguf
|
||||
|
|
@ -1695,6 +1722,7 @@ export function SharedComposer({
|
|||
setText(e.currentTarget.value);
|
||||
}}
|
||||
onKeyDown={onKeyDown}
|
||||
onPaste={handleFilePaste}
|
||||
onBlur={() => {
|
||||
// Mac: switching input methods can fire compositionstart without a
|
||||
// matching compositionend, leaving composingRef pinned. The OS always
|
||||
|
|
|
|||
248
studio/frontend/src/features/chat/utils/clipboard-files.ts
Normal file
248
studio/frontend/src/features/chat/utils/clipboard-files.ts
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
|
||||
const MAX_NATIVE_IMAGE_DIMENSION = 8192;
|
||||
const MAX_NATIVE_IMAGE_RGBA_BYTES = 64 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_BYTES = 20 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_FILES = 8;
|
||||
|
||||
type ClipboardPasteEvent = {
|
||||
readonly clipboardData: DataTransfer | null;
|
||||
readonly defaultPrevented: boolean;
|
||||
readonly isTrusted: boolean;
|
||||
preventDefault: () => void;
|
||||
};
|
||||
type NativeClipboardFile = {
|
||||
readonly name: string;
|
||||
readonly mimeType: string;
|
||||
readonly base64: string;
|
||||
};
|
||||
|
||||
function browserClipboardFiles(clipboardData: DataTransfer): File[] {
|
||||
const files = Array.from(clipboardData.files).filter((file) => file.size > 0);
|
||||
if (files.length > 0) return files;
|
||||
|
||||
return Array.from(clipboardData.items)
|
||||
.filter((item) => item.kind === "file")
|
||||
.map((item) => item.getAsFile())
|
||||
.filter((file): file is File => file !== null && file.size > 0);
|
||||
}
|
||||
|
||||
function clipboardTypes(clipboardData: DataTransfer): string[] {
|
||||
return Array.from(clipboardData.types, (type) => type.toLowerCase());
|
||||
}
|
||||
|
||||
function clipboardHasLocalFileUri(
|
||||
clipboardData: DataTransfer,
|
||||
types: readonly string[],
|
||||
): boolean {
|
||||
const uriTypes = types.filter(
|
||||
(type) => type.includes("uri-list") || type.includes("urilist"),
|
||||
);
|
||||
for (const type of uriTypes) {
|
||||
try {
|
||||
if (
|
||||
clipboardData
|
||||
.getData(type)
|
||||
.split(/\r?\n/)
|
||||
.some((line) => line.trim().toLowerCase().startsWith("file:"))
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function clipboardHasPlainText(clipboardData: DataTransfer): boolean {
|
||||
try {
|
||||
return clipboardData.getData("text/plain").length > 0;
|
||||
} catch {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
function validDimension(value: number): boolean {
|
||||
return (
|
||||
Number.isSafeInteger(value) &&
|
||||
value > 0 &&
|
||||
value <= MAX_NATIVE_IMAGE_DIMENSION
|
||||
);
|
||||
}
|
||||
|
||||
function canvasPng(canvas: HTMLCanvasElement): Promise<Blob | null> {
|
||||
return new Promise((resolve) => canvas.toBlob(resolve, "image/png"));
|
||||
}
|
||||
|
||||
function isLinuxDesktop(): boolean {
|
||||
if (typeof navigator === "undefined") return false;
|
||||
return `${navigator.platform} ${navigator.userAgent}`.toLowerCase().includes("linux");
|
||||
}
|
||||
|
||||
async function readNativeClipboardFiles(): Promise<File[]> {
|
||||
const { invoke } = await import("@tauri-apps/api/core");
|
||||
const nativeFiles = await invoke<NativeClipboardFile[]>(
|
||||
"read_native_clipboard_files",
|
||||
);
|
||||
if (nativeFiles.length > MAX_CLIPBOARD_FILES) return [];
|
||||
|
||||
let totalBytes = 0;
|
||||
const files: File[] = [];
|
||||
for (const file of nativeFiles) {
|
||||
if (
|
||||
!file.name ||
|
||||
file.name.length > 255 ||
|
||||
file.name.includes("/") ||
|
||||
file.name.includes("\0") ||
|
||||
file.base64.length > Math.ceil((MAX_CLIPBOARD_BYTES * 4) / 3) + 4
|
||||
) {
|
||||
return [];
|
||||
}
|
||||
const binary = globalThis.atob(file.base64);
|
||||
const bytes = new Uint8Array(binary.length);
|
||||
for (let index = 0; index < binary.length; index += 1) {
|
||||
bytes[index] = binary.charCodeAt(index);
|
||||
}
|
||||
totalBytes += bytes.byteLength;
|
||||
if (totalBytes > MAX_CLIPBOARD_BYTES) return [];
|
||||
files.push(
|
||||
new File([bytes], file.name, {
|
||||
type: file.mimeType || "application/octet-stream",
|
||||
lastModified: Date.now(),
|
||||
}),
|
||||
);
|
||||
}
|
||||
return files;
|
||||
}
|
||||
|
||||
async function readLinuxClipboardImage(): Promise<File | null> {
|
||||
const { invoke } = await import("@tauri-apps/api/core");
|
||||
const raw = await invoke<ArrayBuffer | Uint8Array>("read_native_clipboard_png");
|
||||
const png = Uint8Array.from(raw instanceof Uint8Array ? raw : new Uint8Array(raw));
|
||||
if (png.byteLength === 0 || png.byteLength > MAX_CLIPBOARD_BYTES) return null;
|
||||
return new File([png], "pasted-image.png", {
|
||||
type: "image/png",
|
||||
lastModified: Date.now(),
|
||||
});
|
||||
}
|
||||
|
||||
async function readNativeClipboardImage(): Promise<File | null> {
|
||||
let image: Awaited<ReturnType<
|
||||
typeof import("@tauri-apps/plugin-clipboard-manager").readImage
|
||||
>> | null = null;
|
||||
|
||||
try {
|
||||
if (isLinuxDesktop()) return await readLinuxClipboardImage();
|
||||
const { readImage } = await import("@tauri-apps/plugin-clipboard-manager");
|
||||
image = await readImage();
|
||||
const { width, height } = await image.size();
|
||||
if (!validDimension(width) || !validDimension(height)) return null;
|
||||
|
||||
const expectedRgbaBytes = width * height * 4;
|
||||
if (expectedRgbaBytes > MAX_NATIVE_IMAGE_RGBA_BYTES) return null;
|
||||
|
||||
const rgba = await image.rgba();
|
||||
if (rgba.byteLength !== expectedRgbaBytes) return null;
|
||||
|
||||
const canvas = document.createElement("canvas");
|
||||
canvas.width = width;
|
||||
canvas.height = height;
|
||||
try {
|
||||
const context = canvas.getContext("2d");
|
||||
if (!context) return null;
|
||||
const pixels = new Uint8ClampedArray(
|
||||
rgba.buffer as ArrayBuffer,
|
||||
rgba.byteOffset,
|
||||
rgba.byteLength,
|
||||
);
|
||||
context.putImageData(new ImageData(pixels, width, height), 0, 0);
|
||||
const blob = await canvasPng(canvas);
|
||||
if (!blob || blob.size === 0 || blob.size > MAX_CLIPBOARD_BYTES) {
|
||||
return null;
|
||||
}
|
||||
return new File([blob], "pasted-image.png", {
|
||||
type: "image/png",
|
||||
lastModified: Date.now(),
|
||||
});
|
||||
} finally {
|
||||
canvas.width = 0;
|
||||
canvas.height = 0;
|
||||
}
|
||||
} catch {
|
||||
return null;
|
||||
} finally {
|
||||
if (image) {
|
||||
try {
|
||||
await image.close();
|
||||
} catch {
|
||||
// The native resource may already have been released after an invoke failure.
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function addClipboardFiles(
|
||||
files: readonly File[],
|
||||
addFiles: (files: readonly File[]) => void | Promise<void>,
|
||||
onError?: () => void,
|
||||
): void {
|
||||
void Promise.resolve(addFiles(files)).catch(() => onError?.());
|
||||
}
|
||||
|
||||
function addNativeClipboardFiles(
|
||||
addFiles: (files: readonly File[]) => void | Promise<void>,
|
||||
onError?: () => void,
|
||||
): void {
|
||||
void (async () => {
|
||||
try {
|
||||
const files = await readNativeClipboardFiles();
|
||||
if (files.length > 0) return files;
|
||||
} catch {
|
||||
// The clipboard may contain image pixels instead of file paths.
|
||||
}
|
||||
const image = await readNativeClipboardImage();
|
||||
return image ? [image] : [];
|
||||
})().then((files) => {
|
||||
if (files.length > 0) addClipboardFiles(files, addFiles, onError);
|
||||
else onError?.();
|
||||
});
|
||||
}
|
||||
|
||||
export function pasteClipboardFiles(
|
||||
event: ClipboardPasteEvent,
|
||||
addFiles: (files: readonly File[]) => void | Promise<void>,
|
||||
onError?: () => void,
|
||||
): void {
|
||||
const { clipboardData } = event;
|
||||
if (clipboardData) {
|
||||
const browserFiles = browserClipboardFiles(clipboardData);
|
||||
if (browserFiles.length > 0) {
|
||||
event.preventDefault();
|
||||
addClipboardFiles(browserFiles, addFiles, onError);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (!isTauri || !event.isTrusted || event.defaultPrevented) return;
|
||||
if (!clipboardData) {
|
||||
addNativeClipboardFiles(addFiles, onError);
|
||||
return;
|
||||
}
|
||||
|
||||
const types = clipboardTypes(clipboardData);
|
||||
const advertisesImage = types.some((type) => type.startsWith("image/"));
|
||||
const advertisesFile =
|
||||
types.includes("files") ||
|
||||
types.some((type) => type.includes("copied-files")) ||
|
||||
clipboardHasLocalFileUri(clipboardData, types);
|
||||
if (!advertisesImage && !advertisesFile && clipboardHasPlainText(clipboardData)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (advertisesImage || advertisesFile) event.preventDefault();
|
||||
addNativeClipboardFiles(addFiles, onError);
|
||||
}
|
||||
|
|
@ -8,6 +8,78 @@ type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
|
|||
const THINK_OPEN_TAG = "<think>";
|
||||
const THINK_CLOSE_TAG = "</think>";
|
||||
|
||||
/**
|
||||
* Normalize streamed string or structured delta content to inline text.
|
||||
* Structured reasoning-only chunks remain distinguishable so their fallback
|
||||
* timer can span consecutive chunks even though each chunk carries closed tags.
|
||||
*/
|
||||
export function extractDeltaText(delta: unknown): {
|
||||
text: string;
|
||||
structuredReasoningContinues: boolean;
|
||||
} {
|
||||
const extractReasoningText = (payload: unknown): string => {
|
||||
if (typeof payload === "string") return payload;
|
||||
if (Array.isArray(payload)) {
|
||||
return payload.map((item) => extractReasoningText(item)).join("");
|
||||
}
|
||||
if (!payload || typeof payload !== "object") return "";
|
||||
|
||||
const obj = payload as Record<string, unknown>;
|
||||
for (const key of ["thinking", "text", "content", "reasoning", "summary"]) {
|
||||
if (key in obj) {
|
||||
const text = extractReasoningText(obj[key]);
|
||||
if (text) return text;
|
||||
}
|
||||
}
|
||||
return "";
|
||||
};
|
||||
|
||||
if (typeof delta === "string") {
|
||||
return { text: delta, structuredReasoningContinues: false };
|
||||
}
|
||||
if (!Array.isArray(delta)) {
|
||||
return { text: "", structuredReasoningContinues: false };
|
||||
}
|
||||
|
||||
let text = "";
|
||||
let structuredReasoningContinues = false;
|
||||
for (const part of delta) {
|
||||
if (typeof part === "string") {
|
||||
text += part;
|
||||
if (part) {
|
||||
structuredReasoningContinues = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (!part || typeof part !== "object") continue;
|
||||
const obj = part as {
|
||||
type?: string;
|
||||
text?: string;
|
||||
content?: string;
|
||||
thinking?: string;
|
||||
};
|
||||
if (obj.type === "text" || obj.type === "output_text") {
|
||||
const visibleText =
|
||||
typeof obj.text === "string"
|
||||
? obj.text
|
||||
: typeof obj.content === "string"
|
||||
? obj.content
|
||||
: "";
|
||||
text += visibleText;
|
||||
if (visibleText) {
|
||||
structuredReasoningContinues = false;
|
||||
}
|
||||
} else if (obj.type === "thinking" || obj.type === "reasoning") {
|
||||
const thinking = extractReasoningText(obj);
|
||||
if (thinking) {
|
||||
text += `${THINK_OPEN_TAG}${thinking}${THINK_CLOSE_TAG}`;
|
||||
structuredReasoningContinues = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return { text, structuredReasoningContinues };
|
||||
}
|
||||
|
||||
// ContentPart from @assistant-ui/react has readonly fields, so coalescing via
|
||||
// `last.text += text` fails (TS2540). Instead replace the last element with a
|
||||
// fresh merged object: same allocation cost as mutation but type-safe.
|
||||
|
|
@ -64,6 +136,6 @@ export function parseAssistantContent(
|
|||
return parts;
|
||||
}
|
||||
|
||||
export function hasClosedThinkTag(raw: string): boolean {
|
||||
return raw.includes(THINK_CLOSE_TAG);
|
||||
export function hasUnclosedThinkTag(raw: string): boolean {
|
||||
return raw.lastIndexOf(THINK_OPEN_TAG) > raw.lastIndexOf(THINK_CLOSE_TAG);
|
||||
}
|
||||
|
|
|
|||
218
studio/frontend/src/features/chat/utils/reasoning-duration.ts
Normal file
218
studio/frontend/src/features/chat/utils/reasoning-duration.ts
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
type MessagePartLike = {
|
||||
type?: unknown;
|
||||
text?: unknown;
|
||||
};
|
||||
|
||||
type ReasoningMetadata = {
|
||||
reasoningDuration?: unknown;
|
||||
reasoningDurations?: unknown;
|
||||
};
|
||||
|
||||
function asDuration(value: unknown): number | undefined {
|
||||
return typeof value === "number" && Number.isFinite(value) && value >= 0
|
||||
? value
|
||||
: undefined;
|
||||
}
|
||||
|
||||
function getReasoningGroupIndex(
|
||||
parts: readonly MessagePartLike[],
|
||||
endIndex: number,
|
||||
): number {
|
||||
let index = -1;
|
||||
let previousWasReasoning = false;
|
||||
|
||||
const limit = Math.min(endIndex, parts.length - 1);
|
||||
for (let partIndex = 0; partIndex <= limit; partIndex += 1) {
|
||||
const isReasoning = parts[partIndex]?.type === "reasoning";
|
||||
if (isReasoning && !previousWasReasoning) {
|
||||
index += 1;
|
||||
}
|
||||
previousWasReasoning = isReasoning;
|
||||
}
|
||||
|
||||
return index;
|
||||
}
|
||||
|
||||
export function countReasoningGroups(
|
||||
parts: readonly MessagePartLike[],
|
||||
): number {
|
||||
return getReasoningGroupIndex(parts, parts.length - 1) + 1;
|
||||
}
|
||||
|
||||
/**
|
||||
* Total reasoning text in the LAST reasoning group (the group any new
|
||||
* reasoning would join). The adapter compares this across chunks to tell "the
|
||||
* model is still thinking" from "the model has moved on to the answer": a
|
||||
* provider that closes every reasoning block atomically would otherwise freeze
|
||||
* the group's timer at its first close.
|
||||
*/
|
||||
export function lastReasoningGroupTextLength(
|
||||
parts: readonly MessagePartLike[],
|
||||
): number {
|
||||
let total = 0;
|
||||
let inGroup = false;
|
||||
for (let index = parts.length - 1; index >= 0; index -= 1) {
|
||||
if (parts[index]?.type !== "reasoning") {
|
||||
if (inGroup) break;
|
||||
continue;
|
||||
}
|
||||
inGroup = true;
|
||||
const text = parts[index]?.text;
|
||||
total += typeof text === "string" ? text.length : 0;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
export function resolveReasoningGroupDuration(
|
||||
parts: readonly MessagePartLike[],
|
||||
startIndex: number,
|
||||
custom: ReasoningMetadata | null | undefined,
|
||||
): number | undefined {
|
||||
const index = getReasoningGroupIndex(parts, startIndex);
|
||||
if (index < 0) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (Array.isArray(custom?.reasoningDurations)) {
|
||||
return asDuration(custom.reasoningDurations[index]);
|
||||
}
|
||||
|
||||
if (index !== getReasoningGroupIndex(parts, parts.length - 1)) {
|
||||
return undefined;
|
||||
}
|
||||
return asDuration(custom?.reasoningDuration);
|
||||
}
|
||||
|
||||
export function createReasoningDurationTracker(
|
||||
now: () => number = Date.now,
|
||||
) {
|
||||
let durations: number[] = [];
|
||||
// First time each group index became visible. A group can be closed and
|
||||
// reopened -- a provider that emits several complete <think>...</think>
|
||||
// blocks in a row has them coalesced into one rendered group -- so the
|
||||
// duration is always measured from the first sighting, not the last.
|
||||
const startedAt: number[] = [];
|
||||
let activeIndex: number | null = null;
|
||||
let groupCount = 0;
|
||||
// Reasoning text seen so far per group, used to decide whether a closed
|
||||
// group is still growing and should reopen.
|
||||
const reasoningLength: number[] = [];
|
||||
// The group a server summary would land on. The backend emits one summary at
|
||||
// the end of each visible reasoning pass, before the next pass can begin, so
|
||||
// "the group that started most recently" is the correct target. (A FIFO queue
|
||||
// is tempting but wrong: it mis-assigns as soon as one group has no summary.)
|
||||
let serverSummaryTargetIndex: number | null = null;
|
||||
// Indices whose duration came from the server; local timing must not
|
||||
// overwrite an authoritative value.
|
||||
const serverClaimed = new Set<number>();
|
||||
|
||||
const setDuration = (index: number, duration: number) => {
|
||||
if (durations[index] === duration) {
|
||||
return;
|
||||
}
|
||||
const next = [...durations];
|
||||
next[index] = duration;
|
||||
durations = next;
|
||||
};
|
||||
const measure = (index: number, finishedAt: number) => {
|
||||
if (serverClaimed.has(index)) {
|
||||
return;
|
||||
}
|
||||
const from = startedAt[index];
|
||||
if (from === undefined) {
|
||||
return;
|
||||
}
|
||||
setDuration(index, Math.max(0, Math.round((finishedAt - from) / 1000)));
|
||||
};
|
||||
const finishGroupAt = (finishedAt: number) => {
|
||||
if (activeIndex === null) {
|
||||
return;
|
||||
}
|
||||
const index = activeIndex;
|
||||
activeIndex = null;
|
||||
measure(index, finishedAt);
|
||||
};
|
||||
|
||||
return {
|
||||
get groupCount() {
|
||||
return groupCount;
|
||||
},
|
||||
get hasActiveGroup() {
|
||||
return activeIndex !== null;
|
||||
},
|
||||
startGroup(index = groupCount) {
|
||||
if (activeIndex === index) {
|
||||
return;
|
||||
}
|
||||
const at = now();
|
||||
finishGroupAt(at);
|
||||
// A single delta can reveal more than one group at once. Any index we
|
||||
// skipped became visible and closed within this same chunk, so give it a
|
||||
// measured zero rather than leaving a hole in the persisted array.
|
||||
for (let skipped = groupCount; skipped < index; skipped += 1) {
|
||||
if (startedAt[skipped] === undefined) {
|
||||
startedAt[skipped] = at;
|
||||
}
|
||||
measure(skipped, at);
|
||||
}
|
||||
if (startedAt[index] === undefined) {
|
||||
startedAt[index] = at;
|
||||
}
|
||||
activeIndex = index;
|
||||
groupCount = Math.max(groupCount, index + 1);
|
||||
serverSummaryTargetIndex = index;
|
||||
},
|
||||
/**
|
||||
* Reopen a group that already closed, but only while its reasoning text is
|
||||
* still growing. Providers that emit each reasoning block as a complete
|
||||
* <think>...</think> chunk close the group on every chunk; without this the
|
||||
* group would freeze at the first close. Gating on growth is what keeps the
|
||||
* timer from running on into the answer.
|
||||
*/
|
||||
resumeGroup(index: number, currentReasoningLength: number) {
|
||||
const seen = reasoningLength[index] ?? 0;
|
||||
if (currentReasoningLength <= seen) {
|
||||
return;
|
||||
}
|
||||
reasoningLength[index] = currentReasoningLength;
|
||||
if (activeIndex === index || startedAt[index] === undefined) {
|
||||
return;
|
||||
}
|
||||
finishGroupAt(now());
|
||||
activeIndex = index;
|
||||
},
|
||||
finishGroup() {
|
||||
finishGroupAt(now());
|
||||
},
|
||||
recordServerDuration(reasoningMs: unknown): boolean {
|
||||
if (
|
||||
typeof reasoningMs !== "number" ||
|
||||
!Number.isFinite(reasoningMs) ||
|
||||
reasoningMs < 0
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
if (serverSummaryTargetIndex !== null) {
|
||||
serverClaimed.add(serverSummaryTargetIndex);
|
||||
setDuration(
|
||||
serverSummaryTargetIndex,
|
||||
Math.max(0, Math.round(reasoningMs / 1000)),
|
||||
);
|
||||
serverSummaryTargetIndex = null;
|
||||
}
|
||||
return true;
|
||||
},
|
||||
metadata() {
|
||||
if (durations.length === 0) {
|
||||
return {};
|
||||
}
|
||||
return {
|
||||
reasoningDuration: durations.at(-1) ?? 0,
|
||||
reasoningDurations: durations,
|
||||
};
|
||||
},
|
||||
};
|
||||
}
|
||||
236
studio/frontend/tests/reasoning-duration.test.ts
Normal file
236
studio/frontend/tests/reasoning-duration.test.ts
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
|
||||
import {
|
||||
countReasoningGroups,
|
||||
createReasoningDurationTracker,
|
||||
lastReasoningGroupTextLength,
|
||||
resolveReasoningGroupDuration,
|
||||
} from "../src/features/chat/utils/reasoning-duration.ts";
|
||||
import { extractDeltaText } from "../src/features/chat/utils/parse-assistant-content.ts";
|
||||
|
||||
const separatedReasoning = [
|
||||
{ type: "reasoning" },
|
||||
{ type: "tool-call" },
|
||||
{ type: "reasoning" },
|
||||
{ type: "text" },
|
||||
];
|
||||
|
||||
test("selects per-group durations while preserving legacy messages", () => {
|
||||
const current = {
|
||||
reasoningDuration: 5,
|
||||
reasoningDurations: [2, 5],
|
||||
};
|
||||
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 0, current), 2);
|
||||
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 2, current), 5);
|
||||
assert.equal(countReasoningGroups(separatedReasoning), 2);
|
||||
|
||||
const legacy = { reasoningDuration: 5 };
|
||||
assert.equal(
|
||||
resolveReasoningGroupDuration(separatedReasoning, 0, legacy),
|
||||
undefined,
|
||||
);
|
||||
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 2, legacy), 5);
|
||||
|
||||
const contiguous = [
|
||||
{ type: "reasoning" },
|
||||
{ type: "reasoning" },
|
||||
{ type: "text" },
|
||||
];
|
||||
assert.equal(countReasoningGroups(contiguous), 1);
|
||||
assert.equal(
|
||||
resolveReasoningGroupDuration(contiguous, 0, {
|
||||
reasoningDurations: [3],
|
||||
}),
|
||||
3,
|
||||
);
|
||||
});
|
||||
|
||||
test("tracks the exact reasoning, tool, reasoning sequence", () => {
|
||||
let now = 0;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
tracker.startGroup();
|
||||
now = 1_200;
|
||||
tracker.recordServerDuration(2_000);
|
||||
tracker.finishGroup();
|
||||
|
||||
tracker.startGroup();
|
||||
now = 5_600;
|
||||
tracker.recordServerDuration(5_000);
|
||||
tracker.finishGroup();
|
||||
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 5,
|
||||
reasoningDurations: [2, 5],
|
||||
});
|
||||
});
|
||||
|
||||
test("keeps groups aligned when summaries are missing or orphaned", () => {
|
||||
let now = 0;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
tracker.startGroup();
|
||||
now = 2_000;
|
||||
tracker.finishGroup();
|
||||
|
||||
tracker.startGroup();
|
||||
now = 7_000;
|
||||
tracker.recordServerDuration(5_000);
|
||||
tracker.finishGroup();
|
||||
tracker.recordServerDuration(9_000);
|
||||
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 5,
|
||||
reasoningDurations: [2, 5],
|
||||
});
|
||||
});
|
||||
|
||||
test("accepts zero after closure and rejects malformed server timing", () => {
|
||||
let now = 0;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
tracker.startGroup();
|
||||
now = 1_000;
|
||||
tracker.finishGroup();
|
||||
assert.equal(tracker.recordServerDuration(0), true);
|
||||
assert.equal(tracker.recordServerDuration(-1), false);
|
||||
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 0,
|
||||
reasoningDurations: [0],
|
||||
});
|
||||
});
|
||||
|
||||
test("omits unknown timing and falls back to elapsed time", () => {
|
||||
let now = 0;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
tracker.startGroup();
|
||||
assert.deepEqual(tracker.metadata(), {});
|
||||
|
||||
now = 3_200;
|
||||
tracker.finishGroup();
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 3,
|
||||
reasoningDurations: [3],
|
||||
});
|
||||
});
|
||||
|
||||
test("keeps structured reasoning active only when it is the final content", () => {
|
||||
assert.deepEqual(
|
||||
extractDeltaText([{ type: "reasoning", text: "First" }]),
|
||||
{
|
||||
text: "<think>First</think>",
|
||||
structuredReasoningContinues: true,
|
||||
},
|
||||
);
|
||||
assert.deepEqual(
|
||||
extractDeltaText([
|
||||
{ type: "reasoning", text: "Last thought" },
|
||||
{ type: "text", text: "Answer" },
|
||||
]),
|
||||
{
|
||||
text: "<think>Last thought</think>Answer",
|
||||
structuredReasoningContinues: false,
|
||||
},
|
||||
);
|
||||
assert.deepEqual(
|
||||
extractDeltaText([
|
||||
{ type: "text", text: "Preface" },
|
||||
{ type: "reasoning", text: "First thought" },
|
||||
]),
|
||||
{
|
||||
text: "Preface<think>First thought</think>",
|
||||
structuredReasoningContinues: true,
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
test("keeps a coalesced reasoning group growing across atomic blocks", () => {
|
||||
let now = 1_770_000_000_000;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
// A provider that closes every reasoning block in its own chunk still
|
||||
// belongs to ONE rendered group, so the timer must span all of them.
|
||||
tracker.startGroup();
|
||||
tracker.resumeGroup(0, "first block".length);
|
||||
tracker.finishGroup();
|
||||
|
||||
now += 3_000;
|
||||
tracker.resumeGroup(0, "first blocksecond block".length);
|
||||
tracker.finishGroup();
|
||||
|
||||
// The answer that follows adds no reasoning text, so the timer stops here.
|
||||
now += 3_000;
|
||||
tracker.resumeGroup(0, "first blocksecond block".length);
|
||||
tracker.finishGroup();
|
||||
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 3,
|
||||
reasoningDurations: [3],
|
||||
});
|
||||
});
|
||||
|
||||
test("never persists a hole when one delta reveals several groups", () => {
|
||||
let now = 1_770_000_000_000;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
// Index 0 was never started explicitly: it became visible and closed inside
|
||||
// the same chunk that revealed index 1.
|
||||
tracker.startGroup(1);
|
||||
now += 4_000;
|
||||
tracker.finishGroup();
|
||||
|
||||
const metadata = tracker.metadata();
|
||||
const durations = metadata.reasoningDurations as number[];
|
||||
assert.equal(durations.length, 2);
|
||||
assert.ok(durations.every((value) => typeof value === "number"));
|
||||
assert.deepEqual(JSON.parse(JSON.stringify(durations)), [0, 4]);
|
||||
});
|
||||
|
||||
test("a server duration is never overwritten by local timing", () => {
|
||||
let now = 1_770_000_000_000;
|
||||
const tracker = createReasoningDurationTracker(() => now);
|
||||
|
||||
tracker.startGroup();
|
||||
tracker.recordServerDuration(2_000);
|
||||
now += 30_000;
|
||||
tracker.resumeGroup(0, 99);
|
||||
tracker.finishGroup();
|
||||
|
||||
assert.deepEqual(tracker.metadata(), {
|
||||
reasoningDuration: 2,
|
||||
reasoningDurations: [2],
|
||||
});
|
||||
});
|
||||
|
||||
test("lastReasoningGroupTextLength measures only the last reasoning group", () => {
|
||||
assert.equal(
|
||||
lastReasoningGroupTextLength([
|
||||
{ type: "reasoning", text: "aaaa" },
|
||||
{ type: "tool-call" },
|
||||
{ type: "reasoning", text: "bb" },
|
||||
{ type: "reasoning", text: "c" },
|
||||
]),
|
||||
3,
|
||||
);
|
||||
// The answer that follows is not reasoning, so it does not count -- but the
|
||||
// group itself is still measured, which is what lets resumeGroup see that the
|
||||
// reasoning has stopped growing.
|
||||
assert.equal(
|
||||
lastReasoningGroupTextLength([
|
||||
{ type: "reasoning", text: "aaaa" },
|
||||
{ type: "text", text: "answer" },
|
||||
]),
|
||||
4,
|
||||
);
|
||||
assert.equal(
|
||||
lastReasoningGroupTextLength([{ type: "text", text: "answer only" }]),
|
||||
0,
|
||||
);
|
||||
assert.equal(lastReasoningGroupTextLength([]), 0);
|
||||
});
|
||||
28
studio/frontend/tsconfig.test.json
Normal file
28
studio/frontend/tsconfig.test.json
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"lib": ["ES2022", "DOM", "DOM.Iterable"],
|
||||
"module": "ESNext",
|
||||
"types": ["node"],
|
||||
"skipLibCheck": true,
|
||||
|
||||
"moduleResolution": "bundler",
|
||||
"allowImportingTsExtensions": true,
|
||||
"verbatimModuleSyntax": true,
|
||||
"moduleDetection": "force",
|
||||
"noEmit": true,
|
||||
"jsx": "react-jsx",
|
||||
|
||||
"strict": true,
|
||||
"noUnusedLocals": false,
|
||||
"noUnusedParameters": true,
|
||||
"erasableSyntaxOnly": true,
|
||||
"noFallthroughCasesInSwitch": true,
|
||||
"noUncheckedSideEffectImports": true,
|
||||
"baseUrl": ".",
|
||||
"paths": {
|
||||
"@/*": ["./src/*"]
|
||||
}
|
||||
},
|
||||
"include": ["tests"]
|
||||
}
|
||||
5
studio/src-tauri/Cargo.lock
generated
5
studio/src-tauri/Cargo.lock
generated
|
|
@ -5566,10 +5566,15 @@ checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
|||
name = "unsloth-studio"
|
||||
version = "2026.4.8"
|
||||
dependencies = [
|
||||
"arboard",
|
||||
"base64 0.22.1",
|
||||
"dirs",
|
||||
"elevated-command",
|
||||
"fix-path-env",
|
||||
"gdk",
|
||||
"gdk-pixbuf",
|
||||
"glib",
|
||||
"gtk",
|
||||
"hmac",
|
||||
"libc",
|
||||
"log",
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ fix-path-env = { git = "https://github.com/tauri-apps/fix-path-env-rs" }
|
|||
tauri-plugin-opener = "2.5.4"
|
||||
tauri-plugin-updater = "2"
|
||||
tauri-plugin-clipboard-manager = "2"
|
||||
arboard = "3.6.1"
|
||||
tauri-plugin-dialog = "2"
|
||||
rand = "0.10.0"
|
||||
tauri-plugin-notification = "2.3.3"
|
||||
|
|
@ -36,6 +37,10 @@ tauri-plugin-window-state = "2"
|
|||
libc = "0.2"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
gdk = "0.18"
|
||||
gdk-pixbuf = "0.18"
|
||||
glib = "0.18"
|
||||
gtk = "0.18"
|
||||
elevated-command = "1.1.2"
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@
|
|||
},
|
||||
"updater:default",
|
||||
"clipboard-manager:allow-write-text",
|
||||
"clipboard-manager:allow-read-image",
|
||||
"window-state:default"
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ mod desktop_update_policy;
|
|||
mod diagnostics;
|
||||
mod install;
|
||||
mod native_backend_lease;
|
||||
mod native_clipboard;
|
||||
mod native_file_dialogs;
|
||||
mod native_intents;
|
||||
mod native_path_policy;
|
||||
|
|
@ -218,6 +219,8 @@ fn main() {
|
|||
desktop_update_policy::check_desktop_manual_update,
|
||||
desktop_update_policy::desktop_update_policy,
|
||||
diagnostics::collect_support_diagnostics,
|
||||
native_clipboard::read_native_clipboard_files,
|
||||
native_clipboard::read_native_clipboard_png,
|
||||
native_file_dialogs::save_native_file,
|
||||
native_file_dialogs::pick_native_chat_import,
|
||||
native_intents::drain_native_intents,
|
||||
|
|
|
|||
442
studio/src-tauri/src/native_clipboard.rs
Normal file
442
studio/src-tauri/src/native_clipboard.rs
Normal file
|
|
@ -0,0 +1,442 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
|
||||
use serde::Serialize;
|
||||
use std::fs::File;
|
||||
use std::io::Read;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const MAX_CLIPBOARD_IMAGE_DIMENSION: i32 = 8192;
|
||||
const MAX_CLIPBOARD_RGBA_BYTES: u64 = 64 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_PNG_BYTES: usize = 20 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_SOURCE_BYTES: u64 = 20 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_TOTAL_BYTES: u64 = 20 * 1024 * 1024;
|
||||
const MAX_CLIPBOARD_FILES: usize = 8;
|
||||
const MAX_CLIPBOARD_CANDIDATES: usize = 32;
|
||||
#[cfg(target_os = "linux")]
|
||||
const MAX_CLIPBOARD_URI_BYTES: usize = 64 * 1024;
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NativeClipboardFile {
|
||||
name: String,
|
||||
mime_type: String,
|
||||
base64: String,
|
||||
}
|
||||
|
||||
fn validate_dimensions(width: i32, height: i32) -> Result<(), String> {
|
||||
if width <= 0
|
||||
|| height <= 0
|
||||
|| width > MAX_CLIPBOARD_IMAGE_DIMENSION
|
||||
|| height > MAX_CLIPBOARD_IMAGE_DIMENSION
|
||||
{
|
||||
return Err("Clipboard image dimensions are invalid or too large.".to_string());
|
||||
}
|
||||
let rgba_bytes = (width as u64)
|
||||
.checked_mul(height as u64)
|
||||
.and_then(|pixels| pixels.checked_mul(4))
|
||||
.ok_or_else(|| "Clipboard image dimensions overflow.".to_string())?;
|
||||
if rgba_bytes > MAX_CLIPBOARD_RGBA_BYTES {
|
||||
return Err("Clipboard image pixel data is too large.".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn validate_png_bytes(png: &[u8]) -> Result<(), String> {
|
||||
const SIGNATURE: &[u8; 8] = b"\x89PNG\r\n\x1a\n";
|
||||
if png.len() < 24 || png.len() > MAX_CLIPBOARD_PNG_BYTES || &png[..8] != SIGNATURE {
|
||||
return Err("Clipboard PNG data is invalid or too large.".to_string());
|
||||
}
|
||||
if &png[12..16] != b"IHDR" {
|
||||
return Err("Clipboard PNG header is invalid.".to_string());
|
||||
}
|
||||
let width = u32::from_be_bytes(png[16..20].try_into().unwrap());
|
||||
let height = u32::from_be_bytes(png[20..24].try_into().unwrap());
|
||||
let width = i32::try_from(width).map_err(|_| "Clipboard PNG width is invalid.".to_string())?;
|
||||
let height =
|
||||
i32::try_from(height).map_err(|_| "Clipboard PNG height is invalid.".to_string())?;
|
||||
validate_dimensions(width, height)
|
||||
}
|
||||
|
||||
fn clipboard_file_mime_type(path: &Path) -> Option<&'static str> {
|
||||
let extension = path
|
||||
.extension()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
let mime_type = match extension.as_str() {
|
||||
"json" | "jsonl" | "ndjson" => "application/json",
|
||||
"md" | "markdown" | "mdx" => "text/markdown",
|
||||
"csv" => "text/csv",
|
||||
"html" | "htm" => "text/html",
|
||||
"xml" => "application/xml",
|
||||
"svg" => "image/svg+xml",
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"png" => "image/png",
|
||||
"webp" => "image/webp",
|
||||
"gif" => "image/gif",
|
||||
"pdf" => "application/pdf",
|
||||
"docx" => "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"odt" => "application/vnd.oasis.opendocument.text",
|
||||
"ods" => "application/vnd.oasis.opendocument.spreadsheet",
|
||||
"mp3" => "audio/mpeg",
|
||||
"wav" => "audio/wav",
|
||||
"m4a" => "audio/mp4",
|
||||
"ogg" | "oga" => "audio/ogg",
|
||||
"flac" => "audio/flac",
|
||||
"aac" => "audio/aac",
|
||||
"txt" | "text" | "log" | "rst" | "tsv" | "yaml" | "yml" | "toml" | "ini" | "cfg"
|
||||
| "conf" | "env" | "properties" | "css" | "scss" | "sass" | "less" | "js" | "jsx"
|
||||
| "mjs" | "cjs" | "ts" | "tsx" | "py" | "pyi" | "ipynb" | "rb" | "php" | "go" | "rs"
|
||||
| "java" | "kt" | "kts" | "scala" | "swift" | "c" | "h" | "cc" | "cpp" | "hpp" | "cxx"
|
||||
| "cs" | "m" | "mm" | "sh" | "bash" | "zsh" | "fish" | "ps1" | "bat" | "lua" | "pl"
|
||||
| "pm" | "r" | "jl" | "dart" | "vue" | "svelte" | "astro" | "sql" | "graphql" | "gql"
|
||||
| "proto" | "tf" | "tfvars" | "gradle" | "dockerfile" | "makefile" | "cmake" | "diff"
|
||||
| "patch" => "text/plain",
|
||||
_ => return None,
|
||||
};
|
||||
Some(mime_type)
|
||||
}
|
||||
|
||||
fn open_regular_clipboard_file(path: &Path) -> Option<File> {
|
||||
let metadata = std::fs::symlink_metadata(path).ok()?;
|
||||
if metadata.file_type().is_symlink() || !metadata.is_file() {
|
||||
return None;
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
std::fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.custom_flags(libc::O_NONBLOCK | libc::O_NOFOLLOW)
|
||||
.open(path)
|
||||
.ok()
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
File::open(path).ok()
|
||||
}
|
||||
}
|
||||
|
||||
fn read_clipboard_files(paths: Vec<PathBuf>) -> Result<Vec<NativeClipboardFile>, String> {
|
||||
let mut remaining = MAX_CLIPBOARD_TOTAL_BYTES;
|
||||
let mut files = Vec::new();
|
||||
for path in paths.into_iter().take(MAX_CLIPBOARD_CANDIDATES) {
|
||||
if remaining == 0 || files.len() >= MAX_CLIPBOARD_FILES {
|
||||
break;
|
||||
}
|
||||
let Some(name) = path
|
||||
.file_name()
|
||||
.map(|value| value.to_string_lossy().into_owned())
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(mime_type) = clipboard_file_mime_type(&path) else {
|
||||
continue;
|
||||
};
|
||||
let Some(source) = open_regular_clipboard_file(&path) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(metadata) = source.metadata() else {
|
||||
continue;
|
||||
};
|
||||
let limit = MAX_CLIPBOARD_SOURCE_BYTES.min(remaining);
|
||||
if !metadata.is_file() || metadata.len() > limit {
|
||||
continue;
|
||||
}
|
||||
let mut bytes = Vec::with_capacity(metadata.len() as usize);
|
||||
if source.take(limit + 1).read_to_end(&mut bytes).is_err() || bytes.len() as u64 > limit {
|
||||
continue;
|
||||
}
|
||||
remaining -= bytes.len() as u64;
|
||||
files.push(NativeClipboardFile {
|
||||
name,
|
||||
mime_type: mime_type.to_string(),
|
||||
base64: BASE64.encode(bytes),
|
||||
});
|
||||
}
|
||||
if files.is_empty() {
|
||||
return Err("Clipboard does not contain readable local files.".to_string());
|
||||
}
|
||||
Ok(files)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn encode_clipboard_pixbuf(image: &gdk_pixbuf::Pixbuf) -> Result<Vec<u8>, String> {
|
||||
validate_dimensions(image.width(), image.height())?;
|
||||
let png = image
|
||||
.save_to_bufferv("png", &[])
|
||||
.map_err(|error| format!("Could not encode clipboard image: {error}"))?;
|
||||
if png.is_empty() || png.len() > MAX_CLIPBOARD_PNG_BYTES {
|
||||
return Err("Clipboard image encoding is empty or too large.".to_string());
|
||||
}
|
||||
Ok(png)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn local_clipboard_path(uri: &str) -> Option<PathBuf> {
|
||||
let (path, hostname) = glib::filename_from_uri(uri).ok()?;
|
||||
hostname.is_none().then_some(path)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn local_clipboard_paths_from_bytes(data: &[u8]) -> Vec<PathBuf> {
|
||||
if data.len() > MAX_CLIPBOARD_URI_BYTES {
|
||||
return Vec::new();
|
||||
}
|
||||
let Ok(text) = std::str::from_utf8(data) else {
|
||||
return Vec::new();
|
||||
};
|
||||
text.lines()
|
||||
.map(|line| {
|
||||
line.trim_matches(|character: char| character.is_whitespace() || character == '\0')
|
||||
})
|
||||
.filter_map(local_clipboard_path)
|
||||
.take(MAX_CLIPBOARD_CANDIDATES)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn read_gtk_clipboard_paths() -> Vec<PathBuf> {
|
||||
let clipboard = gtk::Clipboard::get(&gdk::SELECTION_CLIPBOARD);
|
||||
let mut paths: Vec<PathBuf> = clipboard
|
||||
.wait_for_uris()
|
||||
.into_iter()
|
||||
.filter_map(|uri| local_clipboard_path(uri.as_str()))
|
||||
.take(MAX_CLIPBOARD_CANDIDATES)
|
||||
.collect();
|
||||
|
||||
for target in clipboard.wait_for_targets().unwrap_or_default() {
|
||||
if paths.len() >= MAX_CLIPBOARD_CANDIDATES {
|
||||
break;
|
||||
}
|
||||
if !target.name().to_ascii_lowercase().contains("copied-files") {
|
||||
continue;
|
||||
}
|
||||
let Some(data) = clipboard.wait_for_contents(&target) else {
|
||||
continue;
|
||||
};
|
||||
let length = data.length();
|
||||
if length <= 0 || length as usize > MAX_CLIPBOARD_URI_BYTES {
|
||||
continue;
|
||||
}
|
||||
for path in local_clipboard_paths_from_bytes(&data.data()) {
|
||||
if !paths.contains(&path) {
|
||||
paths.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
paths.truncate(MAX_CLIPBOARD_CANDIDATES);
|
||||
paths
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
async fn native_clipboard_paths() -> Result<Vec<PathBuf>, String> {
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
glib::MainContext::default().invoke(move || {
|
||||
let _ = tx.send(read_gtk_clipboard_paths());
|
||||
});
|
||||
rx.await
|
||||
.map_err(|_| "Clipboard file reader stopped unexpectedly.".to_string())
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
async fn native_clipboard_paths() -> Result<Vec<PathBuf>, String> {
|
||||
tokio::task::spawn_blocking(|| {
|
||||
let mut clipboard = arboard::Clipboard::new().map_err(|error| error.to_string())?;
|
||||
let mut paths = clipboard
|
||||
.get()
|
||||
.file_list()
|
||||
.map_err(|error| error.to_string())?;
|
||||
paths.truncate(MAX_CLIPBOARD_CANDIDATES);
|
||||
Ok(paths)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "Clipboard file reader stopped unexpectedly.".to_string())?
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn read_native_clipboard_files(
|
||||
window: tauri::WebviewWindow,
|
||||
) -> Result<Vec<NativeClipboardFile>, String> {
|
||||
crate::native_intents::ensure_main_window(&window)?;
|
||||
let paths = native_clipboard_paths().await?;
|
||||
tokio::task::spawn_blocking(move || read_clipboard_files(paths))
|
||||
.await
|
||||
.map_err(|_| "Clipboard file loader stopped unexpectedly.".to_string())?
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn read_gtk_clipboard_file_image() -> Result<gdk_pixbuf::Pixbuf, String> {
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
for path in read_gtk_clipboard_paths() {
|
||||
let Some(source) = open_regular_clipboard_file(&path) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(metadata) = source.metadata() else {
|
||||
continue;
|
||||
};
|
||||
if !metadata.is_file() || metadata.len() > MAX_CLIPBOARD_SOURCE_BYTES {
|
||||
continue;
|
||||
}
|
||||
let descriptor_path = PathBuf::from(format!("/proc/self/fd/{}", source.as_raw_fd()));
|
||||
let Some((_, width, height)) = gdk_pixbuf::Pixbuf::file_info(&descriptor_path) else {
|
||||
continue;
|
||||
};
|
||||
if validate_dimensions(width, height).is_err() {
|
||||
continue;
|
||||
}
|
||||
let Ok(image) = gdk_pixbuf::Pixbuf::from_file(&descriptor_path) else {
|
||||
continue;
|
||||
};
|
||||
if validate_dimensions(image.width(), image.height()).is_ok() {
|
||||
return Ok(image);
|
||||
}
|
||||
}
|
||||
Err("Clipboard does not contain a readable image or local image file.".to_string())
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn read_gtk_clipboard_png() -> Result<Vec<u8>, String> {
|
||||
let clipboard = gtk::Clipboard::get(&gdk::SELECTION_CLIPBOARD);
|
||||
let png_target = gdk::Atom::intern("image/png");
|
||||
if let Some(data) = clipboard.wait_for_contents(&png_target) {
|
||||
let length = data.length();
|
||||
if length <= 0 || length as usize > MAX_CLIPBOARD_PNG_BYTES {
|
||||
return Err("Clipboard PNG data is empty or too large.".to_string());
|
||||
}
|
||||
let png = data.data();
|
||||
validate_png_bytes(&png)?;
|
||||
return Ok(png);
|
||||
}
|
||||
|
||||
let image = match clipboard.wait_for_image() {
|
||||
Some(image) => image,
|
||||
None => read_gtk_clipboard_file_image()?,
|
||||
};
|
||||
encode_clipboard_pixbuf(&image)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[tauri::command]
|
||||
pub async fn read_native_clipboard_png(
|
||||
window: tauri::WebviewWindow,
|
||||
) -> Result<tauri::ipc::Response, String> {
|
||||
crate::native_intents::ensure_main_window(&window)?;
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
glib::MainContext::default().invoke(move || {
|
||||
let _ = tx.send(read_gtk_clipboard_png());
|
||||
});
|
||||
let png = rx
|
||||
.await
|
||||
.map_err(|_| "Clipboard image reader stopped unexpectedly.".to_string())??;
|
||||
Ok(tauri::ipc::Response::new(png))
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
#[tauri::command]
|
||||
pub async fn read_native_clipboard_png(
|
||||
window: tauri::WebviewWindow,
|
||||
) -> Result<tauri::ipc::Response, String> {
|
||||
crate::native_intents::ensure_main_window(&window)?;
|
||||
Err("Native PNG clipboard fallback is only available on Linux.".to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn clipboard_dimensions_are_bounded() {
|
||||
assert!(validate_dimensions(3840, 2160).is_ok());
|
||||
assert!(validate_dimensions(0, 100).is_err());
|
||||
assert!(validate_dimensions(8193, 100).is_err());
|
||||
assert!(validate_dimensions(8192, 8192).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clipboard_file_mime_types_cover_text_attachments() {
|
||||
assert_eq!(
|
||||
clipboard_file_mime_type(Path::new("data.json")),
|
||||
Some("application/json")
|
||||
);
|
||||
assert_eq!(
|
||||
clipboard_file_mime_type(Path::new("notes.md")),
|
||||
Some("text/markdown")
|
||||
);
|
||||
assert_eq!(clipboard_file_mime_type(Path::new("unknown.bin")), None);
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[test]
|
||||
fn clipboard_png_headers_are_bounded_before_decode() {
|
||||
let mut png = vec![0; 24];
|
||||
png[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
|
||||
png[12..16].copy_from_slice(b"IHDR");
|
||||
png[16..20].copy_from_slice(&1920_u32.to_be_bytes());
|
||||
png[20..24].copy_from_slice(&1080_u32.to_be_bytes());
|
||||
assert!(validate_png_bytes(&png).is_ok());
|
||||
png[16..20].copy_from_slice(&9000_u32.to_be_bytes());
|
||||
assert!(validate_png_bytes(&png).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clipboard_file_reads_are_bounded() {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let path = directory.path().join("notes.md");
|
||||
std::fs::write(&path, b"clipboard text").unwrap();
|
||||
let oversized = directory.path().join("oversized.md");
|
||||
File::create(&oversized)
|
||||
.unwrap()
|
||||
.set_len(MAX_CLIPBOARD_SOURCE_BYTES + 1)
|
||||
.unwrap();
|
||||
|
||||
let files =
|
||||
read_clipboard_files(vec![directory.path().to_path_buf(), oversized, path]).unwrap();
|
||||
assert_eq!(files.len(), 1);
|
||||
assert_eq!(files[0].name, "notes.md");
|
||||
assert_eq!(files[0].mime_type, "text/markdown");
|
||||
|
||||
assert_eq!(BASE64.decode(&files[0].base64).unwrap(), b"clipboard text");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn clipboard_file_reads_reject_symlinks() {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let target = directory.path().join("target.md");
|
||||
let link = directory.path().join("link.md");
|
||||
std::fs::write(&target, b"clipboard text").unwrap();
|
||||
std::os::unix::fs::symlink(&target, &link).unwrap();
|
||||
|
||||
assert!(open_regular_clipboard_file(&link).is_none());
|
||||
assert!(read_clipboard_files(vec![link]).is_err());
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[test]
|
||||
fn copied_file_targets_parse_local_uris() {
|
||||
let paths = local_clipboard_paths_from_bytes(
|
||||
b"copy\nfile:///tmp/pasted%20notes.md\nhttps://example.com/ignored.md\0",
|
||||
);
|
||||
assert_eq!(paths, vec![PathBuf::from("/tmp/pasted notes.md")]);
|
||||
assert!(
|
||||
local_clipboard_paths_from_bytes(&vec![b'x'; MAX_CLIPBOARD_URI_BYTES + 1]).is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[test]
|
||||
fn clipboard_file_uris_must_be_local() {
|
||||
assert_eq!(
|
||||
local_clipboard_path("file:///tmp/pasted%20image.png"),
|
||||
Some(PathBuf::from("/tmp/pasted image.png"))
|
||||
);
|
||||
assert!(local_clipboard_path("file://remote/tmp/image.png").is_none());
|
||||
assert!(local_clipboard_path("https://example.com/image.png").is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -20,10 +20,15 @@ THREAD_SIDEBAR = FRONTEND / "features/chat/thread-sidebar.tsx"
|
|||
SHARED_COMPOSER = FRONTEND / "features/chat/shared-composer.tsx"
|
||||
TITLEBAR = FRONTEND / "components/tauri/window-titlebar.tsx"
|
||||
NATIVE_DIALOGS = REPO / "studio/src-tauri/src/native_file_dialogs.rs"
|
||||
NATIVE_CLIPBOARD = REPO / "studio/src-tauri/src/native_clipboard.rs"
|
||||
TAURI_MAIN = REPO / "studio/src-tauri/src/main.rs"
|
||||
|
||||
|
||||
APP_PROVIDER = FRONTEND / "app/provider.tsx"
|
||||
|
||||
CLIPBOARD_FILES = FRONTEND / "features/chat/utils/clipboard-files.ts"
|
||||
TAURI_CAPABILITIES = REPO / "studio/src-tauri/capabilities/default.json"
|
||||
|
||||
|
||||
def test_file_actions_route_through_native_commands_only_in_tauri():
|
||||
helper = NATIVE_FILES.read_text(encoding = "utf-8")
|
||||
|
|
@ -98,6 +103,73 @@ def test_chat_exports_await_native_saves_and_markdown_uses_shared_helper():
|
|||
assert "downloadFile(" in thread
|
||||
|
||||
|
||||
def test_clipboard_file_paste_is_bounded_and_wired_to_both_composers():
|
||||
helper = CLIPBOARD_FILES.read_text(encoding = "utf-8")
|
||||
thread = THREAD.read_text(encoding = "utf-8")
|
||||
shared_composer = SHARED_COMPOSER.read_text(encoding = "utf-8")
|
||||
capabilities = TAURI_CAPABILITIES.read_text(encoding = "utf-8")
|
||||
|
||||
for contract in (
|
||||
"clipboardData.files",
|
||||
"clipboardData.items",
|
||||
"item.getAsFile()",
|
||||
"file.size > 0",
|
||||
'clipboardData.getData("text/plain")',
|
||||
"event.isTrusted",
|
||||
"event.defaultPrevented",
|
||||
'types.includes("files")',
|
||||
'type.includes("uri-list")',
|
||||
'"read_native_clipboard_files"',
|
||||
"globalThis.atob(file.base64)",
|
||||
"new File([bytes], file.name",
|
||||
"MAX_CLIPBOARD_BYTES",
|
||||
'import("@tauri-apps/plugin-clipboard-manager")',
|
||||
"await readImage()",
|
||||
"rgba.byteLength !== expectedRgbaBytes",
|
||||
"await image.close()",
|
||||
):
|
||||
assert contract in helper
|
||||
|
||||
assert "addAttachmentOnPaste={false}" in thread
|
||||
assert "onPaste={handleFilePaste}" in thread
|
||||
assert "pasteClipboardFiles" in thread
|
||||
assert "aui.composer().addAttachment(file)" in thread
|
||||
assert "onPaste={handleFilePaste}" in shared_composer
|
||||
assert "pasteClipboardFiles" in shared_composer
|
||||
assert "addFiles(files)" in shared_composer
|
||||
assert capabilities.count('"clipboard-manager:allow-read-image"') == 1
|
||||
assert '"clipboard-manager:allow-read-text"' not in capabilities
|
||||
|
||||
|
||||
def test_native_clipboard_bridge_is_bounded_and_registered():
|
||||
native_clipboard = NATIVE_CLIPBOARD.read_text(encoding = "utf-8")
|
||||
tauri_main = TAURI_MAIN.read_text(encoding = "utf-8")
|
||||
|
||||
for contract in (
|
||||
"MAX_CLIPBOARD_FILES",
|
||||
"MAX_CLIPBOARD_URI_BYTES",
|
||||
"MAX_CLIPBOARD_TOTAL_BYTES",
|
||||
"MAX_CLIPBOARD_SOURCE_BYTES",
|
||||
"MAX_CLIPBOARD_RGBA_BYTES",
|
||||
".take(limit + 1)",
|
||||
".wait_for_uris()",
|
||||
".wait_for_targets()",
|
||||
'contains("copied-files")',
|
||||
"open_regular_clipboard_file(&path)",
|
||||
'"/proc/self/fd/{}"',
|
||||
".wait_for_image()",
|
||||
"glib::filename_from_uri",
|
||||
"glib::MainContext::default().invoke",
|
||||
"arboard::Clipboard::new()",
|
||||
"BASE64.encode(bytes)",
|
||||
"tauri::ipc::Response::new(png)",
|
||||
):
|
||||
assert contract in native_clipboard
|
||||
|
||||
assert "native_clipboard::read_native_clipboard_files" in tauri_main
|
||||
assert "native_clipboard::read_native_clipboard_png" in tauri_main
|
||||
|
||||
|
||||
def test_desktop_startup_waits_for_auth_without_intermediate_handoff():
|
||||
source = APP_PROVIDER.read_text(encoding = "utf-8")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue