From 0666ad2b2b2452668733729e8b54234f5964643a Mon Sep 17 00:00:00 2001 From: Mario Limonciello Date: Mon, 10 Aug 2026 12:53:12 -0500 Subject: [PATCH 001/211] ci : target ROCm 7.14 for build and release (#25775) * Switch ROCm from 7.2.1 to 7.14 ROCm 7.14 is the first production release using TheRock build system. It can be installed using multi-arch deliverables from wheels, debs, rpms, tarballs or runfiles. Adjust ROCm targets for Linux and Windows to use this instead. * ci: switch all other Windows ROCm jobs to ROCm 7.14 wheels Move the shared windows-setup-rocm composite action from the HIP SDK PRO Edition installer to the multi-arch ROCm wheels (rocm[libraries,devel]). The wheel-install logic that previously lived inline in release.yml is now in the shared action, and both build-cache.yml and release.yml call it. Also migrate the build-cuda-windows.yml hip job to the same wheel-based layout (cache path/key, rocm-sdk environment setup, llvm/bin compiler paths) so it keeps working after the action's contract changed; drop its now-unused ROCm 7.2.1 rocWMMA download and stale include path. --- .github/actions/windows-setup-rocm/action.yml | 28 +- .github/workflows/build-cache.yml | 10 +- .github/workflows/build-cuda-windows.yml | 76 +++-- .github/workflows/release.yml | 317 +++++++++--------- 4 files changed, 233 insertions(+), 198 deletions(-) diff --git a/.github/actions/windows-setup-rocm/action.yml b/.github/actions/windows-setup-rocm/action.yml index fd9f8e5a4168..aecbcf14f522 100644 --- a/.github/actions/windows-setup-rocm/action.yml +++ b/.github/actions/windows-setup-rocm/action.yml @@ -8,8 +8,26 @@ inputs: runs: using: "composite" steps: - - name: Setup ROCm - uses: ./.github/actions/install-exe - with: - url: https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ inputs.version }}-Win11-For-HIP.exe - args: -install + - name: Install ROCm with Wheels + shell: pwsh + run: | + $ErrorActionPreference = "Stop" + write-host "Setting up Python virtual environment" + + # Create the venv directly at the cache location to avoid relocation issues + New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null + python -m venv C:\TheRock\build\.venv + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + write-host "Upgrading pip" + python -m pip install --upgrade pip + + write-host "Installing ROCm wheels for multi-arch support" + # Install ROCm wheels for multi-arch support (this may take several minutes) + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ inputs.version }}" + + # Pre-expand the devel tree so it is included in the cache + write-host "Initializing ROCm devel tree" + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + write-host "Completed ROCm wheel installation to C:\TheRock\build" diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 327f71978bf1..e15fc5e0830f 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -123,8 +123,8 @@ jobs: runs-on: windows-2022 env: - # Make sure this is in sync with build.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + # Make sure this is in sync with release.yml and build-cuda-windows.yml + ROCM_VERSION: "7.14.0" steps: - name: Clone @@ -135,11 +135,11 @@ jobs: uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} diff --git a/.github/workflows/build-cuda-windows.yml b/.github/workflows/build-cuda-windows.yml index 367a3a8546cf..ff900802f4a2 100644 --- a/.github/workflows/build-cuda-windows.yml +++ b/.github/workflows/build-cuda-windows.yml @@ -83,7 +83,7 @@ jobs: env: # Make sure this is in sync with build-cache.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + ROCM_VERSION: "7.14.0" strategy: matrix: @@ -97,36 +97,53 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Use ROCm Installation Cache + - name: Cache ROCm Installation uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH - name: Verify ROCm id: verify run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version + # Test the ROCm clang shipped in the installed wheel + & "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -134,28 +151,27 @@ jobs: # TODO: this build does not match the build in release.yml, so we use a different cache key # ideally, the builds should match, similar to the CUDA build above so that we would be able # to populate the ccache for the release with manual runs of this workflow - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} - name: Build id: cmake_build run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/" ` + -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` + -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` + -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_BUILD_TYPE=Release ` -DLLAMA_BUILD_BORINGSSL=ON ` - -DROCM_DIR="${env:HIP_PATH}" ` + -DHIP_PATH="${env:HIP_PATH}" ` -DGGML_HIP=ON ` - -DGPU_TARGETS="gfx1100" ` + -DGPU_TARGETS="gfx1100" ` -DGGML_RPC=ON cmake --build build -j ${env:NUMBER_OF_PROCESSORS} - name: ccache-clear uses: ./.github/actions/ccache-clear with: - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c668930b0caa..3a48a57c1b68 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -748,6 +748,132 @@ jobs: path: llama-bin-win-cpu-${{ matrix.arch }}.zip name: llama-bin-win-cpu-${{ matrix.arch }}.zip + windows-rocm: + runs-on: windows-2022 + + strategy: + matrix: + include: + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201" + build: x64 + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: ccache + uses: ggml-org/ccache-action@v1.2.21 + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + evict-old-files: 1d + + - name: Cache ROCm Installation + id: cache-rocm + uses: actions/cache@v5 + with: + path: C:\TheRock\build + key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }} + + - name: Setup ROCm + if: steps.cache-rocm.outputs.cache-hit != 'true' + uses: ./.github/actions/windows-setup-rocm + with: + version: ${{ matrix.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + write-host "CMake path: $cmakePath" + write-host "Bin path: $binPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH + + - name: Build + run: | + mkdir build + cd build + cmake .. ` + -G "Unix Makefiles" ` + -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` + -DCMAKE_BUILD_TYPE=Release ` + -DGGML_BACKEND_DL=ON ` + -DGGML_NATIVE=OFF ` + -DGGML_CPU=ON ` + -DGGML_CPU_ALL_VARIANTS=ON ` + -DGGML_HIP=ON ` + -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` + -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" ` + -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DHIP_PATH="${env:HIP_PATH}" ` + -DGGML_HIP_ROCWMMA_FATTN=ON ` + -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}" + cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS} + + - name: ccache-clear + uses: ./.github/actions/ccache-clear + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + + - name: Verify HIP backend was built + run: | + $hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue + if (-not $hipDll) { + Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build." + Write-Host "Contents of build\bin:" + Get-ChildItem build\bin | Format-Table -AutoSize + exit 1 + } + Write-Host "HIP backend artifact found:" + $hipDll | Format-Table FullName, Length -AutoSize + + - name: Determine tag name + id: tag + uses: ./.github/actions/get-tag-name + + - name: Get ROCm short version + run: | + $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.') + echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV + + - name: Pack artifacts + run: | + cp "LICENSE" "build\bin\" + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\* + + - name: Upload artifacts + uses: actions/upload-artifact@v6 + with: + path: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + name: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + windows: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1168,8 +1294,8 @@ jobs: strategy: matrix: include: - - ROCM_VERSION: "7.2.1" - gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201" + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201" build: 'x64' steps: @@ -1201,38 +1327,36 @@ jobs: run: | sudo apt install -y build-essential git cmake wget - - name: Setup Legacy ROCm - if: matrix.ROCM_VERSION == '7.2.1' - id: legacy_env - run: | - sudo mkdir --parents --mode=0755 /etc/apt/keyrings - wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \ - gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null - - sudo tee /etc/apt/sources.list.d/rocm.list << EOF - deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} jammy main - EOF - - sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF - Package: * - Pin: release o=repo.radeon.com - Pin-Priority: 600 - EOF - - sudo apt update - sudo apt-get install -y libssl-dev rocm-hip-sdk - - - name: Setup TheRock - if: matrix.ROCM_VERSION != '7.2.1' + - name: Setup TheRock with Wheels id: therock_env run: | - wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz - mkdir install - tar -xf *.tar.gz -C install - export ROCM_PATH=$(pwd)/install - echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV - echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV - echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV + # Create Python virtual environment + python3 -m venv .venv + source .venv/bin/activate + + # Install ROCm wheels for build + # libraries = HIP runtime and CMake configs needed for linking + # devel = compilers, headers, static libs + python -m pip install --upgrade pip + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}" + + # Get ROCm installation paths using the rocm-sdk CLI tool + ROCM_PATH=$(rocm-sdk path --root) + CMAKE_PATH=$(rocm-sdk path --cmake) + BIN_PATH=$(rocm-sdk path --bin) + echo "ROCM_PATH=$ROCM_PATH" + echo "CMAKE_PATH=$CMAKE_PATH" + echo "BIN_PATH=$BIN_PATH" + + # Set environment variables + echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV + echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV + echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV + + # Keep venv activated for subsequent steps + echo "$(pwd)/.venv/bin" >> $GITHUB_PATH - name: Build with native CMake HIP support id: cmake_build @@ -1276,129 +1400,6 @@ jobs: path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz - windows-hip: - needs: [check-release, get-version] - if: ${{ needs.check-release.outputs.should_release == 'true' }} - - runs-on: windows-2022 - - permissions: - actions: write - - env: - HIPSDK_INSTALLER_VERSION: "26.Q1" - - strategy: - matrix: - include: - - name: "radeon" - gpu_targets: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - - name: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Cache ROCm Installation - id: cache-rocm - uses: actions/cache@v5 - with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} - - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Install ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' - id: depends - run: | - $ErrorActionPreference = "Stop" - write-host "Downloading AMD HIP SDK Installer" - Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe" - write-host "Installing AMD HIP SDK" - $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru - $completed = $proc.WaitForExit(600000) - if (-not $completed) { - Write-Error "ROCm installation timed out after 10 minutes. Killing the process" - $proc.Kill() - exit 1 - } - if ($proc.ExitCode -ne 0) { - Write-Error "ROCm installation failed with exit code $($proc.ExitCode)" - exit 1 - } - write-host "Completed AMD HIP SDK installation" - - - name: Verify ROCm - id: verify - run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version - - - name: Build - id: cmake_build - run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" - cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/ -Wno-ignored-attributes -Wno-nested-anon-types" ` - -DCMAKE_BUILD_TYPE=Release ` - -DGGML_BACKEND_DL=ON ` - -DGGML_NATIVE=OFF ` - -DGGML_CPU=OFF ` - -DGPU_TARGETS="${{ matrix.gpu_targets }}" ` - -DGGML_HIP=ON ` - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} ` - -DLLAMA_BUILD_BORINGSSL=ON - cmake --build build --target ggml-hip -j ${env:NUMBER_OF_PROCESSORS} - md "build\bin\rocblas\library\" - md "build\bin\hipblaslt\library" - cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\" - cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\" - - - name: ccache-clear - uses: ./.github/actions/ccache-clear - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Pack artifacts - id: pack_artifacts - run: | - 7z a -snl llama-bin-win-hip-${{ matrix.name }}-x64.zip .\build\bin\* - - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - path: llama-bin-win-hip-${{ matrix.name }}-x64.zip - name: llama-bin-win-hip-${{ matrix.name }}-x64.zip - ios-xcode: needs: [check-release, get-version] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1572,7 +1573,7 @@ jobs: - windows-cpu - windows-cuda #- windows-sycl - - windows-hip + - windows-rocm - windows-openvino - ubuntu-22-rocm - ubuntu-cpu @@ -1684,7 +1685,7 @@ jobs: - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz) - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz) - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz) - - [Ubuntu x64 (ROCm 7.2)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.2-x64.tar.gz) + - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz) - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz) @@ -1702,7 +1703,7 @@ jobs: - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip) - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip) - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip) - - [Windows x64 (HIP)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-hip-radeon-x64.zip) + - [Windows x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-7.14-x64.zip) **openEuler:** - [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23705) From 689e227db485c6b33d061555e74034c93a867649 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Mon, 10 Aug 2026 11:09:19 -0700 Subject: [PATCH 002/211] opencl: transpose the K tile in local memory for FA prefill kernels (#26428) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 18 +++++++ .../ggml-opencl/kernels/flash_attn_f32_f16.cl | 51 ++++++++++++++++-- .../kernels/flash_attn_f32_q4_0.cl | 47 ++++++++++++++-- .../kernels/flash_attn_f32_q8_0.cl | 53 +++++++++++++++++-- 4 files changed, 156 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index fc0fce0d780a..9874ffe8d21e 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)( //------------------------------------------------------------------------------ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor); + static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor); static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor); static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); @@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) { opts += " -D FA_C8_NO_SG_PIN"; } + // Transposed K tile in local memory: the KV rows the QK loop walks together become + // adjacent, so a group of them is ONE 128-bit local read instead of several narrow + // ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a + // but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on + // fa=1 prefill. Output is bit-identical -- only the layout moves. + // + // DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across + // rounds; padding the row stride does not recover it, so the cause is not a simple bank + // conflict and the wider tile does not want this layout. + // + // Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile. + { + const char * e = getenv("GGML_OPENCL_FA_K_LDS_T"); + if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) { + opts += " -D FA_K_LDS_T"; + } + } return opts; } diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl index 6e43ee81e73b..bf7695a2c1d7 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl @@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); +#ifdef FA_K_LDS_T + // K tile transposed: [dk vec][kv row] instead of [kv row][dk vec]. + // + // The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major + // those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they + // are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes, + // no extra registers, arithmetic untouched. + // + // This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS + // read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept + // every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op). + // Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4, + // and the element type only obliges the compiler to align this array to 8. The indices + // are even so the offset is a multiple of 16, but the base has to be too, and relying + // on the compiler to over-align it is relying on luck. + __local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16))); +#define FA_LK(ROW, C) l_k[C][ROW] + // Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and + // BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base. +#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J])) +#else __local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC]; +#define FA_LK(ROW, C) l_k[ROW][C] +#endif __local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC]; #if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE) @@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME( #ifdef FA_K_IMG if (use_kv_pad) { const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; } else { const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row; - l_k[row][col] = read_imageh(k_img, k_row_px + col); + FA_LK(row, col) = read_imageh(k_img, k_row_px + col); } #else const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; #endif } else { - l_k[row][col] = (KV_DATA_TYPE4)(0.0h); + FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h); } } for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) { @@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 2 KV rows adjacent in the transposed tile: one 128-bit local read. + const half8 kk = FA_LK_PAIR(dk_off + k, j); + ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo); + ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi); +#else ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]); ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]); +#endif partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3; partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3; } @@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME( ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { - dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc); + dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc); } local_partial[j][tid] = dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3; @@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 4 KV rows adjacent in the transposed tile: two 128-bit local reads + // instead of four 64-bit ones. + const half8 kk01 = FA_LK_PAIR(k, j); + const half8 kk23 = FA_LK_PAIR(k, j + 2); + dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3); +#else dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0); dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1); dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2); dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3); +#endif } ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl index 95d215971e00..48adba4f725b 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl @@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each + // (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK + // loop is LDS-read-issue-bound. + __local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0( const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; uint k_packed[8]; pack_q4_0_nibbles(qs, k_packed); #pragma unroll for (int j = 0; j < 8; ++j) { - l_k_packed[row][blk * 8 + j] = k_packed[j]; + FA_K_PACKED(row, blk * 8 + j) = k_packed[j]; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0( for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#ifdef FA_K_LDS_T + // 4 KV rows are adjacent in the transposed tile: one 128-bit local + // read per (block, group) instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; const int q_sum = q_sum_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0; + s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1; + s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2; + s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3; +#else s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b]; s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b]; s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b]; s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl index 7e89ed0bd8f1..f50912d21101 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl @@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g]. + // + // The QK loop walks 4 KV rows at a time against the same (b, g), so in the original + // layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local + // reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS + // issues for the same bytes and no extra registers. That matters because the QK loop + // is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS + // reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK + // outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads. + __local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0( const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; #pragma unroll for (int j = 0; j < 8; ++j) { uint k_packed = @@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0( ((uint) qs[j*4 + 1]) << 8 | ((uint) qs[j*4 + 2]) << 16 | ((uint) qs[j*4 + 3]) << 24; - l_k_packed[row][blk * 8 + j] = k_packed; + FA_K_PACKED(row, blk * 8 + j) = k_packed; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0( for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#if defined(FA_K_LDS_T) + // The 4 KV rows are adjacent in the transposed tile, so each (b, g) + // step is ONE 128-bit local read instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)sum0 * qd * ks4.s0; + s1 += (float)sum1 * qd * ks4.s1; + s2 += (float)sum2 * qd * ks4.s2; + s3 += (float)sum3 * qd * ks4.s3; +#else s0 += (float)sum0 * qd * l_k_scale[j ][b]; s1 += (float)sum1 * qd * l_k_scale[j+1][b]; s2 += (float)sum2 * qd * l_k_scale[j+2][b]; s3 += (float)sum3 * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); From 030ebb558a5820b444a8f836ed5cdd46c9b4bd7a Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Tue, 11 Aug 2026 00:02:25 +0530 Subject: [PATCH 003/211] Address review comment of PR 25532 (#26852) --- include/llama.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/include/llama.h b/include/llama.h index c14eb6f50d9a..5a0b66dc25f6 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1324,8 +1324,6 @@ extern "C" { LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p); LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl); LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl); - // copy mutable sampler state without changing dst or its sampling graph bindings - // src and dst must have the same type and configuration LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst); // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); From 84f712946729f8517c972da4eb80db810ffe3210 Mon Sep 17 00:00:00 2001 From: Masashi Yoshimura Date: Tue, 11 Aug 2026 13:10:00 +0900 Subject: [PATCH 004/211] ggml-webgpu: fix CI errors from #25025 and #25262 (#26566) * test new flash_attn test * rebase and fix to disable subgrou matrices when max_kv_tile == 0 * delete log output * Add i32 support to cpy and enables the all ops test * restore the non target ci tests * comment out of TODO of build-cpu.yml * fix format --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 22 +++++++++++++++++-- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 5 ++--- ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl | 2 ++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index 35a55ecaf644..0604e1c2b87b 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -2815,11 +2815,25 @@ class ggml_webgpu_shader_lib { key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16; key.use_sg_matrix = decisions.use_sg_matrix; - const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u, key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.k_direct || key.common.v_direct); - GGML_ASSERT(max_kv_tile > 0); + + // WorkGroup storage size isn't enough for some params with subgroup matrices path (ref. https://github.com/ggml-org/llama.cpp/pull/26566) + if (max_kv_tile == 0) { + GGML_ASSERT(decisions.use_sg_matrix); + // switch to flash_attn_reg_tile path + decisions.use_sg_matrix = false; + decisions.q_tile = GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; + key.common.k_direct = false; + key.common.v_direct = false; + key.use_sg_matrix = false; + max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + context.wg_mem_limit_bytes, decisions.q_tile, 1u, key.common.head_dim_qk, key.common.head_dim_v, + key.common.has_mask, key.common.k_direct || key.common.v_direct); + GGML_ASSERT(max_kv_tile > 0); + } decisions.kv_tile = decisions.use_sg_matrix ? std::min(max_kv_tile, context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES) : @@ -2993,6 +3007,10 @@ class ggml_webgpu_shader_lib { defines.push_back("SRC_F16"); variant += "_f16"; break; + case GGML_TYPE_I32: + defines.push_back("SRC_I32"); + variant += "_i32"; + break; default: GGML_ABORT("Unsupported src type for cpy shader"); } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index ba4b91695fae..98c7162478f8 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4283,9 +4283,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_CPY: case GGML_OP_CONT: - supports_op = ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16)) || - (op->type == GGML_TYPE_I32 && src0->type == GGML_TYPE_F32); + supports_op = (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_I32) && + (src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32); break; case GGML_OP_SET: supports_op = src0->type == src1->type && src0->type == op->type && diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl index 67f1dc0928f8..0d0d81ab650f 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl @@ -4,6 +4,8 @@ enable f16; #define SRC_TYPE f32 #elif defined(SRC_F16) #define SRC_TYPE f16 +#elif defined(SRC_I32) +#define SRC_TYPE i32 #endif #ifdef DST_F32 From 48d22e295e2b86b47366c16390794f3e05ba970a Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Mon, 10 Aug 2026 23:10:31 -0500 Subject: [PATCH 005/211] common/peg : suppress incomplete escape sequences (#26780) --- common/peg-parser.cpp | 19 +++++++++++++++---- tests/peg-parser/test-json-parser.cpp | 24 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index ef290ed7c057..4a4be7cf789f 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -570,23 +570,34 @@ struct parser_executor { } static common_peg_parse_result handle_escape_sequence(common_peg_parse_context & ctx, size_t start, size_t & pos, const char delimiter) { + auto save = pos; + ++pos; // consume '\' if (pos >= ctx.input.size()) { if (!ctx.is_lenient()) { return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + pos = save; // suppress unmatched '\' return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); } char c = ctx.input[pos]; + if (c == delimiter || c == '\\' || c == '/' || c == 'b' || c == 'f' || c == 'n' || c == 'r' || c == 't') { ++pos; return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); - } else if (c == 'u') { - return handle_unicode_escape(ctx, start, pos); - } else { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + + if (c == 'u') { + auto result = handle_unicode_escape(ctx, start, pos); + if (result.need_more_input()) { + pos = save; // suppress incomplete sequence + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); + } + return result; + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } static common_peg_parse_result handle_unicode_escape(common_peg_parse_context & ctx, size_t start, size_t & pos) { diff --git a/tests/peg-parser/test-json-parser.cpp b/tests/peg-parser/test-json-parser.cpp index 5dd00115cead..ec7c2e668ff1 100644 --- a/tests/peg-parser/test-json-parser.cpp +++ b/tests/peg-parser/test-json-parser.cpp @@ -77,6 +77,30 @@ void test_json_parser(testing &t) { t.assert_equal("result_is_need_more_input", true, result.need_more_input()); }); + // Test need_more_input() parsing - incomplete escape sequence in a string value + t.test("need_more_input() parsing - incomplete escape sequence", [](testing &t) { + auto json = build_peg_parser([](common_peg_parser_builder & p) { return p.json(); }); + + std::vector inputs { + R"({"text": "hello\)", // dangling backslash + R"({"text": "hello\u)", // incomplete unicode escape sequence + R"({"text": "hello\u00)", + }; + + for (const auto & input : inputs) { + t.test(input, [&](testing &t) { + common_peg_parse_context ctx(input, COMMON_PEG_PARSE_FLAG_LENIENT); + + auto result = json.parse(ctx); + + t.assert_equal("result_is_need_more_input", true, result.need_more_input()); + + // the incomplete escape sequence is not part of the partial value + t.assert_equal("result_end", input.find('\\'), result.end); + }); + } + }); + t.test("object member", [](testing &t) { auto parser = build_peg_parser([](common_peg_parser_builder & p) { return p.json_member("name", "\"" + p.chars("[a-z]") + "\""); From 14e78ddef7a2061e7d5a31dce4eb7ee0bcdbc840 Mon Sep 17 00:00:00 2001 From: Junmo Kim Date: Tue, 11 Aug 2026 13:20:17 +0900 Subject: [PATCH 006/211] model : fix SWA not being enabled for EXAONE 4.5 (#26848) * model : fix SWA not being enabled for EXAONE 4.5 load_arch_hparams tests `hparams.n_layer() == 64` before LLM_KV_NEXTN_PREDICT_LAYERS has been read. n_layer() returns n_layer_all - n_layer_nextn and n_layer_nextn defaults to 0, so a GGUF carrying the MTP head (block_count=65, nextn=1) evaluates to 65 and the whole SWA block is skipped. The model type switch further down in the same function reads 64, because by then the key has been loaded. n_swa is still filled in by the unconditional get_key below the block, so llama_model_n_swa() reports 4096 and the logs look correct while only swa_type stays LLAMA_SWA_TYPE_NONE. This affects the official LGAI-EXAONE GGUF release as well. EXAONE 4.0 has no MTP head, so block_count is 64 there and the check matches. * model-loader : skip TENSOR_SKIP tensors in the metadata-only path create_tensor asserts on a null buffer type when building from metadata alone, but buft_for_tensor returns null by design for tensors marked TENSOR_SKIP, which is how architectures with nextn/MTP layers mark theirs. Those models cannot be constructed by llama_model_init_from_user at all. The file-backed path below already returns nullptr for the same tensors, so callers see the same thing either way. * tests : cover exaone4 hparams ordering Builds a synthetic exaone4 model with the layout the shipped EXAONE 4.5 GGUFs use (block_count 65 + nextn 1). The swa_type check is the one that catches the ordering bug; the n_layer_nextn and n_layer() checks only tell a broken fixture apart from a real regression. Fails before the ordering fix with "swa_type is not STANDARD", passes after. * Revert "tests : cover exaone4 hparams ordering" This reverts commit d2f3bafeee591ad691396b2708de4baef3aaf602. * Revert "model-loader : skip TENSOR_SKIP tensors in the metadata-only path" This reverts commit aecb9bc0c7896b52afbc43921a1f572aa7b5e53c. --- src/models/exaone4.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/models/exaone4.cpp b/src/models/exaone4.cpp index 863268abcef2..a06819a67caa 100644 --- a/src/models/exaone4.cpp +++ b/src/models/exaone4.cpp @@ -1,6 +1,9 @@ #include "models.h" void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); + if (hparams.n_layer() == 64) { // 32B hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; hparams.n_swa = 4096; @@ -15,9 +18,6 @@ void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); - ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); - - GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); switch (hparams.n_layer()) { case 30: type = LLM_TYPE_1_2B; break; From 4801e3c567d5131dd41b387df5f2d4b1370d92be Mon Sep 17 00:00:00 2001 From: Jim Wu Date: Mon, 10 Aug 2026 21:21:32 -0700 Subject: [PATCH 007/211] tests : disable backend sampler hip multi output (#26878) * test-backend-sampler: skip multi_output_sampling_chain on HIP The new multi_output_sampling_chain test uses top_k, whose backend probs path needs CUB (unavailable on HIP), so sampled_probs is null and the test aborts. Add it to the existing HIP skip list alongside the other TOP_K tests. * ci: keep gpu-rocm logs in a per-run dir keyed by GitHub run id The self-hosted gpu-rocm runner can't upload logs to Azure blob (egress firewalled), so a run's logs were wiped by the next run. Write each run's logs to $OUT/run--/ so an Actions run URL maps to its logs. * test-backend-sampler: also skip multi_output_cpu on HIP Like the other TOP_K-based subtests, multi_output_cpu's backend sampler never initializes on HIP (no CUB TOP_K), so it aborts. Add it to the skip list. --------- Co-authored-by: Jim Wu --- ci/run.sh | 8 ++++++++ tests/test-backend-sampler.cpp | 4 +++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/ci/run.sh b/ci/run.sh index f6c7eb0d5c70..8046df255154 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -49,6 +49,14 @@ mkdir -p "$2" OUT=$(realpath "$1") MNT=$(realpath "$2") +# gpu-rocm self-hosted runner can't upload logs to blob; keep each run's logs in +# their own dir keyed by the GitHub run id so an Actions run URL maps to its logs. +if [ -n "${GG_BUILD_ROCM}" ] && [ -n "${GITHUB_RUN_ID}" ]; then + OUT="$OUT/run-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT:-1}" + mkdir -p "$OUT" + echo "ci results dir: $OUT" +fi + rm -f $OUT/*.log rm -f $OUT/*.exit rm -f $OUT/*.md diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 4b3d12635a36..c23e7248d5e5 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -2105,7 +2105,9 @@ static std::vector collect_tests_to_run(const std::st #ifdef GGML_USE_HIP // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged if (test.name == "penalties" || test.name == "set_sampler" || - test.name == "mixed" || test.name == "top_p") { + test.name == "mixed" || test.name == "top_p" || + test.name == "multi_output_sampling_chain" || + test.name == "multi_output_cpu") { fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str()); continue; } From b3df57286c336255fd02c2162b0b5fe3adc2965e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 11 Aug 2026 09:07:13 +0300 Subject: [PATCH 008/211] tests : clean-up server test, use `tests.sh` in ci (#26886) * tests : remove fetch_server_test_models.py * ci : use tests.sh wrapper of pytest --- .github/workflows/server-sanitize.yml | 4 +- .github/workflows/server-self-hosted.yml | 18 ++-- .github/workflows/server.yml | 12 +-- scripts/fetch_server_test_models.py | 105 ----------------------- tools/server/tests/tests.sh | 5 -- 5 files changed, 17 insertions(+), 127 deletions(-) delete mode 100755 scripts/fetch_server_test_models.py diff --git a/.github/workflows/server-sanitize.yml b/.github/workflows/server-sanitize.yml index 0eeefdf88469..5d696282c709 100644 --- a/.github/workflows/server-sanitize.yml +++ b/.github/workflows/server-sanitize.yml @@ -110,7 +110,7 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -119,4 +119,4 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh diff --git a/.github/workflows/server-self-hosted.yml b/.github/workflows/server-self-hosted.yml index 249f389ff3f2..675ddbaaa580 100644 --- a/.github/workflows/server-self-hosted.yml +++ b/.github/workflows/server-self-hosted.yml @@ -72,7 +72,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -81,7 +81,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -90,7 +90,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -99,7 +99,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-cuda: runs-on: [self-hosted, llama-server, Linux, NVIDIA] @@ -132,7 +132,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -141,7 +141,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -150,7 +150,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -159,7 +159,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-kleidiai: runs-on: ah-ubuntu_22_04-c8g_8x @@ -219,4 +219,4 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index 5a02cc15ad5e..d5abf1d23668 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -104,21 +104,21 @@ jobs: id: server_integration_tests run: | cd tools/server/tests - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow if: ${{ github.event.schedule || github.event.inputs.slow_tests == 'true' }} run: | cd tools/server/tests - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh - name: Tests (Backend sampling) id: server_integration_tests_backend_sampling run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests (Backend sampling) id: server_integration_tests_slow_backend_sampling @@ -126,7 +126,7 @@ jobs: run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh windows: runs-on: windows-2025 @@ -170,7 +170,7 @@ jobs: run: | cd tools/server/tests $env:PYTHONIOENCODING = ":replace" - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -178,4 +178,4 @@ jobs: run: | cd tools/server/tests $env:SLOW_TESTS = "1" - pytest -v -x + ./tests.sh diff --git a/scripts/fetch_server_test_models.py b/scripts/fetch_server_test_models.py deleted file mode 100755 index f43d1f63cdc7..000000000000 --- a/scripts/fetch_server_test_models.py +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env python -''' - This script fetches all the models used in the server tests. - - This is useful for slow tests that use larger models, to avoid them timing out on the model downloads. - - It is meant to be run from the root of the repository. - - Example: - python scripts/fetch_server_test_models.py - ( cd tools/server/tests && ./tests.sh -v -x -m slow ) -''' -import ast -import glob -import logging -import os -from typing import Generator -from pydantic import BaseModel -from typing import Optional -import subprocess - - -class HuggingFaceModel(BaseModel): - hf_repo: str - hf_file: Optional[str] = None - - class Config: - frozen = True - - -def collect_hf_model_test_parameters(test_file) -> Generator[HuggingFaceModel, None, None]: - try: - with open(test_file) as f: - tree = ast.parse(f.read()) - except Exception as e: - logging.error(f'collect_hf_model_test_parameters failed on {test_file}: {e}') - return - - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - for dec in node.decorator_list: - if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Attribute) and dec.func.attr == 'parametrize': - param_names = ast.literal_eval(dec.args[0]).split(",") - if "hf_repo" not in param_names: - continue - - raw_param_values = dec.args[1] - if not isinstance(raw_param_values, ast.List): - logging.warning(f'Skipping non-list parametrize entry at {test_file}:{node.lineno}') - continue - - hf_repo_idx = param_names.index("hf_repo") - hf_file_idx = param_names.index("hf_file") if "hf_file" in param_names else None - - for t in raw_param_values.elts: - if not isinstance(t, ast.Tuple): - logging.warning(f'Skipping non-tuple parametrize entry at {test_file}:{node.lineno}') - continue - yield HuggingFaceModel( - hf_repo=ast.literal_eval(t.elts[hf_repo_idx]), - hf_file=ast.literal_eval(t.elts[hf_file_idx]) if hf_file_idx is not None else None) - - -if __name__ == '__main__': - logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') - - models = sorted(list(set([ - model - for test_file in glob.glob('tools/server/tests/unit/test_*.py') - for model in collect_hf_model_test_parameters(test_file) - ])), key=lambda m: (m.hf_repo, m.hf_file)) - - logging.info(f'Found {len(models)} models in parameterized tests:') - for m in models: - logging.info(f' - {m.hf_repo} / {m.hf_file}') - - cli_path = os.environ.get( - 'LLAMA_CLI_BIN_PATH', - os.path.join( - os.path.dirname(__file__), - '../build/bin/Release/llama-cli.exe' if os.name == 'nt' else '../build/bin/llama-cli')) - - for m in models: - if '<' in m.hf_repo or (m.hf_file is not None and '<' in m.hf_file): - continue - if m.hf_file is not None and '-of-' in m.hf_file: - logging.warning(f'Skipping model at {m.hf_repo} / {m.hf_file} because it is a split file') - continue - logging.info(f'Using llama-cli to ensure model {m.hf_repo}/{m.hf_file} was fetched') - cmd = [ - cli_path, - '-hfr', m.hf_repo, - *([] if m.hf_file is None else ['-hff', m.hf_file]), - '-n', '1', - '-p', 'Hey', - '--no-warmup', - '--log-disable', - '-st'] - if m.hf_file != 'tinyllamas/stories260K.gguf' and 'Mistral-Nemo' not in m.hf_repo: - cmd += ('-fa', 'on') - try: - subprocess.check_call(cmd) - except subprocess.CalledProcessError: - logging.error(f'Failed to fetch model at {m.hf_repo} / {m.hf_file} with command:\n {" ".join(cmd)}') - exit(1) diff --git a/tools/server/tests/tests.sh b/tools/server/tests/tests.sh index 8d6681193db3..433dc99828e4 100755 --- a/tools/server/tests/tests.sh +++ b/tools/server/tests/tests.sh @@ -6,11 +6,6 @@ cd $SCRIPT_DIR set -eu -if [[ "${SLOW_TESTS:-0}" == 1 ]]; then - # Slow tests for tool calls need quite a few models ahead of time to avoid timing out. - python $SCRIPT_DIR/../../../scripts/fetch_server_test_models.py -fi - if [ $# -lt 1 ] then if [[ "${SLOW_TESTS:-0}" == 1 ]]; then From 153d324bcf86d220b235ca010eeb11213f32b5d1 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Tue, 11 Aug 2026 08:20:46 +0200 Subject: [PATCH 009/211] llama: add default load-mode auto, which avoids mmap on iGPUs (#26081) * llama: add new default load-mode auto which picks mmap unless a non-Metal iGPU is used * Update ggml/src/ggml-hexagon/ggml-hexagon.cpp Co-authored-by: Max Krasnyansky * set mmap_support to false on OpenCL backend * fix order of load modes * use -1 for auto * resolve load mode auto earlier to correctly pick gpu host or cpu memory * add load mode auto to llama-bench * bump virtgpu api version, regenerate docs --------- Co-authored-by: Piotr Wilkin (ilintar) Co-authored-by: Max Krasnyansky Co-authored-by: Georgi Gerganov --- common/arg.cpp | 6 ++++-- common/common.h | 2 +- ggml/include/ggml-backend.h | 2 ++ ggml/src/ggml-backend-meta.cpp | 2 ++ ggml/src/ggml-blas/ggml-blas.cpp | 1 + ggml/src/ggml-cann/ggml-cann.cpp | 1 + ggml/src/ggml-cpu/ggml-cpu.cpp | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 1 + ggml/src/ggml-et/ggml-et.cpp | 1 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-opencl/ggml-opencl.cpp | 1 + ggml/src/ggml-openvino/ggml-openvino.cpp | 1 + ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + .../backend/backend-dispatched-device.cpp | 1 + .../backend/shared/api_remoting.h | 2 +- .../ggml-virtgpu/ggml-backend-buffer-type.cpp | 4 ++-- ggml/src/ggml-virtgpu/ggml-backend-device.cpp | 2 +- .../ggml-virtgpu/virtgpu-forward-device.cpp | 4 +++- ggml/src/ggml-virtgpu/virtgpu-forward.gen.h | 3 ++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 + ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + ggml/src/ggml-zdnn/ggml-zdnn.cpp | 3 ++- ggml/src/ggml-zendnn/ggml-zendnn.cpp | 3 ++- include/llama.h | 11 ++++++----- src/llama-model-loader.cpp | 2 +- src/llama-model.cpp | 19 +++++++++++++++++-- src/llama.cpp | 11 +++++++---- tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/llama-bench/llama-bench.cpp | 8 +++++--- tools/server/README.md | 2 +- 33 files changed, 75 insertions(+), 29 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index c37d5cd0aa4b..cb314eee7044 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2605,14 +2605,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_DIO")); add_opt(common_arg( {"-lm", "--load-mode"}, "MODE", - "model loading mode (default: mmap)\n" + "model loading mode (default: auto)\n" + "- auto: mmap, unless a device does not support it\n" "- none: no special loading mode\n" "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n" "- mlock: force system to keep model in RAM rather than swapping or compressing\n" "- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n" "- dio: use DirectIO if available\n", [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } + /**/ if (value == "auto") { params.load_mode = LLAMA_LOAD_MODE_AUTO; } + else if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; } else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; } else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } diff --git a/common/common.h b/common/common.h index 878534dccf50..d485d4fb41f7 100644 --- a/common/common.h +++ b/common/common.h @@ -473,7 +473,7 @@ struct common_params { std::vector fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs - enum llama_load_mode load_mode = LLAMA_LOAD_MODE_MMAP; // how to load the model + enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model common_cpu_params cpuparams; common_cpu_params cpuparams_batch; diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 2924fdbe9884..cc3f8cd36e35 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -154,6 +154,8 @@ extern "C" { bool buffer_from_host_ptr; // event synchronization bool events; + // mmap is supported for loading + bool mmap_support; }; // all the device properties diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad054..7654ea1f30f6 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -132,6 +132,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ false, // Not implemented. /* .buffer_from_host_ptr = */ false, // Not implemented. /* .events = */ false, // Not implemented. + /* .mmap_support = */ true, }; for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { ggml_backend_dev_props tmp_props; @@ -140,6 +141,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; props->caps.events = props->caps.events && tmp_props.caps.events; + props->caps.mmap_support = props->caps.mmap_support && tmp_props.caps.mmap_support; } } diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp index 9745fa29f5db..e4b5bd254747 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp @@ -367,6 +367,7 @@ static void ggml_backend_blas_device_get_props(ggml_backend_dev_t dev, struct gg /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5f51ea3bb3c8..ffa361af4e8b 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2815,6 +2815,7 @@ static void ggml_backend_cann_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 16cc5116c545..c0c9aa3cf09c 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -397,6 +397,7 @@ static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 05e8d7f73684..1d4f4dfbd214 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4801,6 +4801,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ props->type != GGML_BACKEND_DEVICE_TYPE_IGPU, }; } diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b30209095672..e8482f734621 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1646,6 +1646,7 @@ static void ggml_backend_et_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index bdb8af0820a3..f80c60a500bd 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3930,6 +3930,7 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ (bool) opt_hostbuf, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3acff8..ef3c92f27127 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -681,6 +681,7 @@ static void ggml_backend_metal_device_get_props(ggml_backend_dev_t dev, ggml_bac /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 9874ffe8d21e..19dca4ced0b8 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -10795,6 +10795,7 @@ static void ggml_backend_opencl_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index 0e7501fefe38..dfe80e6c86e7 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -763,6 +763,7 @@ static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_ /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 17c53a5f049e..e9de0d0aa98a 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1881,6 +1881,7 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 18d58782ebff..8a19f648bf07 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5649,6 +5649,7 @@ static void ggml_backend_sycl_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp index c7acb8b51ce7..87872df1c7bc 100644 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp +++ b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp @@ -111,6 +111,7 @@ uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_ apir_encode_bool_t(enc, &props.caps.host_buffer); apir_encode_bool_t(enc, &props.caps.buffer_from_host_ptr); apir_encode_bool_t(enc, &props.caps.events); + apir_encode_bool_t(enc, &props.caps.mmap_support); return 0; } diff --git a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h index 6bf97e8a3a24..a5ef3ea476d0 100644 --- a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h +++ b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h @@ -7,7 +7,7 @@ #include #define APIR_PROTOCOL_MAJOR 0 -#define APIR_PROTOCOL_MINOR 1 +#define APIR_PROTOCOL_MINOR 2 #define APIR_HANDSHAKE_MAGIC 0xab1e diff --git a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp index 8fa20ff43bd5..d5bdc993b464 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp @@ -11,9 +11,9 @@ static ggml_backend_buffer_t ggml_backend_remoting_buffer_type_alloc_buffer(ggml context->gpu = gpu; - bool async__unused, host_buffer__unused, events__unused; + bool async__unused, host_buffer__unused, events__unused, mmap_support__unused; bool buffer_from_host_ptr; - apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused); + apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused, &mmap_support__unused); if (buffer_from_host_ptr) { context->apir_context = apir_device_buffer_from_ptr(gpu, size, size); diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp index a978812cd908..987ce9dd110c 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp @@ -65,7 +65,7 @@ static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_ virtgpu * gpu = DEV_TO_GPU(dev); apir_device_get_props(gpu, &props->caps.async, &props->caps.host_buffer, &props->caps.buffer_from_host_ptr, - &props->caps.events); + &props->caps.events, &props->caps.mmap_support); props->caps.buffer_from_host_ptr = false; props->caps.async = false; diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp index 9f513c138dd2..864264f213b5 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp +++ b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp @@ -144,7 +144,8 @@ void apir_device_get_props(virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events) { + bool * events, + bool * mmap_support) { apir_encoder * encoder; apir_decoder * decoder; ApirForwardReturnCode ret; @@ -157,6 +158,7 @@ void apir_device_get_props(virtgpu * gpu, apir_decode_bool_t(decoder, host_buffer); apir_decode_bool_t(decoder, buffer_from_host_ptr); apir_decode_bool_t(decoder, events); + apir_decode_bool_t(decoder, mmap_support); remote_call_finish(gpu, encoder, decoder); diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h index 44b0ad1ffa1d..da28aa5f9047 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h +++ b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h @@ -13,7 +13,8 @@ void apir_device_get_props(struct virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events); + bool * events, + bool * mmap_support); apir_buffer_context_t apir_device_buffer_from_ptr(struct virtgpu * gpu, size_t size, size_t max_tensor_size); /* buffer-type */ diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a923755f9edd..45fa97f81297 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17891,6 +17891,7 @@ static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ true, /* .buffer_from_host_ptr = */ false, /* .events = */ true, + /* .mmap_support = */ !ctx->is_integrated_gpu, }; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 98c7162478f8..6741752b3611 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -3942,6 +3942,7 @@ static void ggml_backend_webgpu_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp index 639b818d128e..4007ac9dfc7d 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp @@ -487,7 +487,8 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp index e6a9b51b7925..ec7ce233145a 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp @@ -654,7 +654,8 @@ static void ggml_backend_zendnn_device_get_props(ggml_backend_dev_t dev, struct /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/include/llama.h b/include/llama.h index 5a0b66dc25f6..bfef0e1d1129 100644 --- a/include/llama.h +++ b/include/llama.h @@ -203,11 +203,12 @@ extern "C" { }; enum llama_load_mode { - LLAMA_LOAD_MODE_NONE = 0, // no special loading mode - LLAMA_LOAD_MODE_MMAP = 1, // memory map the model - LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available + LLAMA_LOAD_MODE_AUTO = -1, // auto-detect based on device capabilities + LLAMA_LOAD_MODE_NONE = 0, // no special loading mode + LLAMA_LOAD_MODE_MMAP = 1, // memory map the model + LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available }; LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 3d50f8a1cb1c..51ba05439682 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -543,7 +543,7 @@ llama_model_loader::llama_model_loader( tensor_buft_overrides = param_tensor_buft_overrides_p; - this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; + this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK || load_mode == LLAMA_LOAD_MODE_AUTO; this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO; if (!fname.empty()) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3bf3a22f26af..cc1917b6da77 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1271,8 +1271,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { this->ml = &ml; // to be used by create_tensor() and load_arch_tensors() + if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO) { + for (const auto & dev : devices) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev.dev, &props); + if (!props.caps.mmap_support) { + ml.use_mmap = false; + break; + } + } + } + + const char * load_mode_name = params.load_mode == LLAMA_LOAD_MODE_AUTO + ? llama_load_mode_name(ml.use_mmap ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE) + : llama_load_mode_name(params.load_mode); + LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n", - __func__, llama_load_mode_name(params.load_mode)); + __func__, load_mode_name); // build a list of buffer types for the CPU and GPU devices pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host); @@ -2452,7 +2467,7 @@ llama_model_params llama_model_default_params() { /*.tensor_buft_overrides =*/ nullptr, /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, - /*.load_mode =*/ LLAMA_LOAD_MODE_MMAP, + /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, diff --git a/src/llama.cpp b/src/llama.cpp index d6e0bbfefa72..94c8f60e0c43 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -48,6 +48,8 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty const char * llama_load_mode_name(enum llama_load_mode load_mode) { switch (load_mode) { + case LLAMA_LOAD_MODE_AUTO: + return "auto"; case LLAMA_LOAD_MODE_NONE: return "none"; case LLAMA_LOAD_MODE_MMAP: @@ -63,11 +65,12 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { } enum llama_load_mode llama_load_mode_from_str(const char * str) { - if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } - if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } - if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } + if (std::strcmp(str, "auto") == 0) { return LLAMA_LOAD_MODE_AUTO; } + if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } + if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } + if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } - if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } throw std::invalid_argument(std::string("unknown load mode: ") + str); } diff --git a/tools/cli/README.md b/tools/cli/README.md index b42b2e5343ff..880c4a540838 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -58,7 +58,7 @@ | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index 552a0c6abf4e..c2e52ac066e7 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -141,7 +141,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c17a27b54019..7c495afe20ef 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -384,7 +384,7 @@ static const cmd_params cmd_params_defaults = { /* n_gpu_layers */ { -1 }, /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, - /* load_mode */ { LLAMA_LOAD_MODE_MMAP }, + /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -459,7 +459,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); - printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -764,7 +764,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { std::vector modes; for (const auto & m : p) { llama_load_mode mode; - if (m == "none") { + if (m == "auto") { + mode = LLAMA_LOAD_MODE_AUTO; + } else if (m == "none") { mode = LLAMA_LOAD_MODE_NONE; } else if (m == "mmap") { mode = LLAMA_LOAD_MODE_MMAP; diff --git a/tools/server/README.md b/tools/server/README.md index 6927caddbb6c..a2ab872b4c1a 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -75,7 +75,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | From 9afff1b7483ddc5a78a7738025b135e0f013d50d Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 11 Aug 2026 12:07:15 +0300 Subject: [PATCH 010/211] tests : fix running server tests on windows (#26889) --- .github/workflows/server.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index d5abf1d23668..9fb4b4ba102f 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -167,15 +167,17 @@ jobs: - name: Tests id: server_integration_tests + shell: bash run: | cd tools/server/tests - $env:PYTHONIOENCODING = ":replace" + export PYTHONIOENCODING=":replace" ./tests.sh - name: Slow tests id: server_integration_tests_slow if: ${{ github.event.schedule || github.event.inputs.slow_tests == 'true' }} + shell: bash run: | cd tools/server/tests - $env:SLOW_TESTS = "1" + export SLOW_TESTS="1" ./tests.sh From 1138b851fae633e9b2e74db0dac3623b7c6fac43 Mon Sep 17 00:00:00 2001 From: Daniel Bevenius Date: Tue, 11 Aug 2026 11:41:38 +0200 Subject: [PATCH 011/211] model-conversion : use save_output_data for causual embeddings [no ci] (#26890) This commit updates the python script that runs the original model to generate embeddings for the causal model, to use save_output_data which stores the token ids and the prompt in addition to logits. The motivation for this is that the embedding logits verification will fail as it expects these files (-prompt.txt and -tokens.bin) to exist. With the changes in this commit the causal-verify-embeddings target works again. --- .../causal/run-casual-gen-embeddings-org.py | 23 +++++-------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py b/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py index b94bec4e765c..cb840dd5504b 100755 --- a/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py +++ b/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py @@ -2,12 +2,15 @@ import argparse import os +import sys import importlib import torch import numpy as np from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM -from pathlib import Path + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) +from utils.common import save_output_data unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME') @@ -54,6 +57,7 @@ prompt = "Hello world today" input_ids = tokenizer(prompt, return_tensors="pt").input_ids # ty: ignore[call-non-callable] +token_ids = input_ids[0].cpu().tolist() print(f"Input tokens: {input_ids}") print(f"Input text: {repr(prompt)}") print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}") # ty: ignore[unresolved-attribute] @@ -74,21 +78,8 @@ print(f"Hidden dimension: {token_embeddings.shape[-1]}") print(f"Number of tokens: {token_embeddings.shape[0]}") - # Save raw token embeddings - data_dir = Path("data") - data_dir.mkdir(exist_ok=True) - bin_filename = data_dir / f"pytorch-{model_name}-embeddings.bin" - txt_filename = data_dir / f"pytorch-{model_name}-embeddings.txt" - - # Save all token embeddings as binary print(token_embeddings) - token_embeddings.astype(np.float32).tofile(bin_filename) - - # Save as text for inspection - with open(txt_filename, "w") as f: - for i, embedding in enumerate(token_embeddings): - for j, val in enumerate(embedding): - f.write(f"{i} {j} {val:.6f}\n") + save_output_data(token_embeddings, token_ids, prompt, model_name, type_suffix="-embeddings") # Print embeddings per token in the requested format print("\nToken embeddings:") @@ -110,5 +101,3 @@ for i, token in enumerate(tokens): print(f" Token {i}: {repr(token)}") - print(f"Saved bin logits to: {bin_filename}") - print(f"Saved txt logist to: {txt_filename}") From 704485942ab54bbbbf1f241b3550ffba35f5f37e Mon Sep 17 00:00:00 2001 From: uvos Date: Tue, 11 Aug 2026 11:47:57 +0200 Subject: [PATCH 012/211] ci: hip-quality-check: update vgpr spill ignore list (#26859) Most of the old ones have been resolved (yay) but the recent refactor of mmq paramters has caused some symbol names to change, leaving a couple of non-ignored failures --- scripts/hip/gcn-cdna-vgpr-check.py | 86 ++---------------------------- 1 file changed, 3 insertions(+), 83 deletions(-) diff --git a/scripts/hip/gcn-cdna-vgpr-check.py b/scripts/hip/gcn-cdna-vgpr-check.py index bbbce52ef39f..40fb789417c3 100644 --- a/scripts/hip/gcn-cdna-vgpr-check.py +++ b/scripts/hip/gcn-cdna-vgpr-check.py @@ -60,90 +60,10 @@ def main(): log_file = sys.argv[1] ignored = { '_ZL21gated_linear_attn_f32ILi128EEviiiifPKfS1_S1_S1_S1_Pf', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', '_ZL13rwkv_wkv7_f32ILi128EEviiiiPKfS1_S1_S1_S1_S1_S1_Pf', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi2ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi64ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL24mul_mat_q_stream_k_fixupIL9ggml_type22ELi8ELb1EEvPKiS2_PfPKfiiimimimi', - '_ZL9mul_mat_qIL9ggml_type3ELi32ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi48ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type20ELi32ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type17ELi64ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL15flash_attn_tileILi256ELi256ELi32ELi1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type19ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type17ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type22ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type19ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type19ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type11ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type11ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL24mul_mat_q_stream_k_fixupIL9ggml_type11ELi128ELb0EEvPKiS2_PfPKfiiimimimi', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type2ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi32ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_vecILi128ELi2EL9ggml_type2ELS0_2ELb0EEvPKcS2_S2_S2_S2_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type10ELi16ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type12ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii' + '_ZL12rwkv_wkv_f32ILi128EEviiiiPKfS1_S1_S1_S1_S1_Pf', + '_ZL9mul_mat_qIL9ggml_type10ELi64ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_', + '_ZL9mul_mat_qIL9ggml_type42ELi128ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_', } functions = parse_log_file(log_file) From 8d274dd7c6233ed73c7509cc2a8be9960f7df7d5 Mon Sep 17 00:00:00 2001 From: Tom Tan <29201606+intel00000@users.noreply.github.com> Date: Tue, 11 Aug 2026 04:42:48 -0700 Subject: [PATCH 013/211] ui: fix context gauge for single-model usage (#25738) * webui: hide loaded model in context gauge at single-model mode * webui: keep context gauge details open state across reopens --- .../ChatFormContextGauge/ContextGaugeDetails.svelte | 9 +++++---- tools/ui/src/lib/hooks/use-context-gauge.svelte.ts | 2 +- tools/ui/src/lib/stores/context-gauge-popup.svelte.ts | 2 +- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte index ab36e4a43f74..eaaba69de622 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte @@ -3,6 +3,7 @@ import { ChevronDown } from '@lucide/svelte'; import * as Collapsible from '$lib/components/ui/collapsible'; import { STATS_UNITS } from '$lib/constants'; + import { gaugePopup } from '$lib/stores/context-gauge-popup.svelte'; interface Props { currentRead: number; @@ -30,19 +31,19 @@ transientDetails }: Props = $props(); - let open = $state(false); - const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0); const hasCurrent = $derived(currentRead > 0 || currentOutput > 0); - + Token usage details - + diff --git a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts index 76dcd356ec29..7ab3f83df609 100644 --- a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts +++ b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts @@ -112,7 +112,7 @@ export function useContextGauge(): UseContextGaugeReturn { return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]); }); const isActiveModelLoaded = $derived( - activeModelId !== null && modelsStore.isModelLoaded(activeModelId) + activeModelId !== null && (!isRouterMode() || modelsStore.isModelLoaded(activeModelId)) ); const isActiveModelLoading = $derived( activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId) diff --git a/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts b/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts index 654dc9f7787b..34f25283f89e 100644 --- a/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts +++ b/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts @@ -16,7 +16,7 @@ import { let closeTimer: ReturnType | undefined; let lastPointerType = ''; -export const gaugePopup = $state({ bottom: 0, centerX: 0, open: false }); +export const gaugePopup = $state({ bottom: 0, centerX: 0, detailsOpen: false, open: false }); function openFrom(trigger: HTMLElement): void { clearTimeout(closeTimer); From 6e62ba538478202094edc6c100c782719e310aa3 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Tue, 11 Aug 2026 14:18:30 +0200 Subject: [PATCH 014/211] mtmd: support pocket-tts (#26871) * adapt the api * text model ok * working impl, need verify and clean up * mtmd: build the pocket-tts transposed convolutions as GEMM + col2im ggml_conv_transpose_1d has no grouped mode, so the depthwise upsample was built as one convolution and one concat per channel, which floods the graph with small nodes and makes kernel launches dominate the decoder. Fold both cases into the column form the seanet decoder already needs: the general case reshapes the kernel to [IC, K * OC] and matmuls it with the input, the depthwise case batches a matmul over the channels so a step scales its own kernel. A single col2im_1d then scatter-adds the columns back to the signal, with the same shape as before, so the overlap-add tail, the streaming state and the bias are untouched. Generation time per frame drops by 80% on CUDA and by 50% on CPU. The output matches the previous implementation sample for sample, with a correlation of 0.999994 and identical frame counts. * flow_temp + frames_after_eos * chunking * mtmd: carry the remaining pocket-tts per-pack settings The language packs also tune the end-of-speech padding and the padding of short prompts, next to the temperature already carried in the mmproj: french_24l asks for 8 tail frames instead of the guessed 3, english_2026-01 asks for short prompts to be padded with spaces. Write both in the mmproj as clip.gen.audio.frames_after_eos and clip.gen.audio.pad_short_text, keyed on the pack in the conversion script like the temperature. The loader keeps them optional, so a mmproj without them behaves as before. Map semicolons to commas for every pack instead, the reference only asks for it on three of them and it costs nothing elsewhere. Existing mmproj files must be converted again to carry the two keys. On a long french text the port now lands within 2% of the reference: 22.96s against 23.44s, with the same peak level and the same amount of silence. * clip.gen.audio.model_variant * clean up code comments * nit: drop the dead flow_temp hparam, the pack table holds the default * update docs * address security problems * less invasive base.py * lint * add mtmd_gen_inp_default * add docs * rm gen_flow_temp --------- Co-authored-by: Pascal --- conversion/__init__.py | 2 + conversion/base.py | 28 ++ conversion/pockettts.py | 378 ++++++++++++++++ gguf-py/gguf/constants.py | 112 +++++ gguf-py/gguf/gguf_writer.py | 3 + src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-model.cpp | 3 + src/models/models.h | 13 + src/models/pockettts.cpp | 146 +++++++ tools/mtmd/CMakeLists.txt | 3 + tools/mtmd/README-dev.md | 15 +- tools/mtmd/clip-impl.h | 40 +- tools/mtmd/clip-model.h | 89 ++++ tools/mtmd/clip.cpp | 303 +++++++++++-- tools/mtmd/clip.h | 5 + tools/mtmd/models/models.h | 56 +++ tools/mtmd/models/pockettts-gen.cpp | 291 ++++++++++++ tools/mtmd/models/pockettts-seanet.cpp | 162 +++++++ tools/mtmd/models/pockettts-spkenc.cpp | 77 ++++ tools/mtmd/models/qwen3tts-gen.cpp | 4 + tools/mtmd/mtmd-audio.cpp | 38 ++ tools/mtmd/mtmd-audio.h | 7 + tools/mtmd/mtmd-helper-gen.cpp | 583 ++++++++++++++++++++++++- tools/mtmd/mtmd-helper.h | 16 +- tools/mtmd/mtmd.cpp | 95 +++- tools/mtmd/mtmd.h | 22 +- tools/tts/README.md | 25 ++ tools/tts/tts.cpp | 15 +- 29 files changed, 2459 insertions(+), 74 deletions(-) create mode 100644 conversion/pockettts.py create mode 100644 src/models/pockettts.cpp create mode 100644 tools/mtmd/models/pockettts-gen.cpp create mode 100644 tools/mtmd/models/pockettts-seanet.cpp create mode 100644 tools/mtmd/models/pockettts-spkenc.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index c7d8046c495f..695289b73ace 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -214,6 +214,7 @@ "Qwen3MoeForCausalLM": "qwen", "Qwen3NextForCausalLM": "qwen", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", @@ -310,6 +311,7 @@ "Qwen2_5_VLForConditionalGeneration": "qwenvl", "Qwen3ASRForConditionalGeneration": "qwen3vl", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", diff --git a/conversion/base.py b/conversion/base.py index a7cd3fd904aa..3572b77c21e2 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -58,6 +58,11 @@ AnyModel = TypeVar("AnyModel", bound="type[ModelBase]") +# for checkpoints that ship no config.json, we will try to provide a synthetic one +HparamsMatcher = Callable[[Path], bool] +HparamsLoader = Callable[[Path], dict[str, Any]] + + class SentencePieceTokenTypes(IntEnum): NORMAL = 1 UNKNOWN = 2 @@ -77,6 +82,7 @@ class ModelBase: ModelType.TEXT: {}, ModelType.MMPROJ: {}, } + _hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = [] dir_model: Path ftype: gguf.LlamaFileType @@ -1040,6 +1046,24 @@ def get_model_part_names(dir_model: Path, prefix: str, suffix: str) -> list[str] return part_names + @staticmethod + def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None: + # some models ship no config.json, will try to guess them + from conversion import load_all_models + load_all_models() + + for matcher, loader in ModelBase._hparams_loaders: + if matcher(dir_model): + return loader(dir_model) + return None + + @classmethod + def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]: + def inner(loader: HparamsLoader) -> HparamsLoader: + cls._hparams_loaders.append((matcher, loader)) + return loader + return inner + @staticmethod def load_hparams(dir_model: Path, is_mistral_format: bool): if is_mistral_format: @@ -1053,6 +1077,10 @@ def load_hparams(dir_model: Path, is_mistral_format: bool): config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict() except Exception as e: logger.warning(f"Failed to load model config from {dir_model}: {e}") + if not (dir_model / "config.json").is_file(): + config = ModelBase.load_hparams_guess(dir_model) + if config is not None: + return config logger.warning("Trying to load config.json instead") with open(dir_model / "config.json", "r", encoding="utf-8") as f: config = json.load(f) diff --git a/conversion/pockettts.py b/conversion/pockettts.py new file mode 100644 index 000000000000..62ecb5acde74 --- /dev/null +++ b/conversion/pockettts.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger + +# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one +# continuous 32-d latent per frame. There is no codebook in this model. +# The checkpoint ships no config.json, hparams come from _load_hparams() below. +# +# Tricks being used to support this model via existing llama.cpp code paths: +# - bos_before_voice and bos_emb are learned input vectors, not tokens +# they are appended to the embedding table as extra tokens, to be looked up like any other row +# - bos_emb lives in latent space, so input_linear is folded into it here +# - the backbone has no lm_head, the embedding table is reused as output for the unused logits +# +# pipeline stage mapping: +# mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder +# flow_lm.transformer --> mapped to normal libllama text model (autoregressive) +# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE +# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV + +# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder +_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731 +_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731 +_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731 +_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731 + +_N_SEANET_STAGES = 3 +_SAMPLE_RATE = 24000 + + +def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]: + part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") + if len(part_names) != 1: + return {} + with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: + return {name: tuple(part[name].shape) for name in part.keys()} + + +@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model)) +def _load_hparams(dir_model: Path) -> dict[str, Any]: + logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") + shapes = _tensor_shapes(dir_model) + n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] + n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] + return { + "architectures": ["PocketTTSModel"], + "model_type": "pockettts", + "num_hidden_layers": n_layer, + "hidden_size": n_embd, + "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], + # the transformer is fully causal with no context limit, this only bounds the KV cache + "max_position_embeddings": 4096, + # not in the checkpoint, but every released variant uses head_dim 64 + "num_attention_heads": n_embd // 64, + # extra rows for the learned input vectors, see _embd_table() + "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), + "rope_theta": 10000.0, + "layer_norm_eps": 1e-5, + "audio_config": { + "num_hidden_layers": n_layer_a, + "hidden_size": n_embd_a, + "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], + "num_attention_heads": n_embd_a // 64, + }, + } + + +@ModelBase.register("PocketTTSModel") +class PocketTTSModel(TextModel): + model_arch = gguf.MODEL_ARCH.POCKETTTS + + _LAYER_TENSOR_MAP = { + "norm1": gguf.MODEL_TENSOR.ATTN_NORM, + "norm2": gguf.MODEL_TENSOR.FFN_NORM, + "self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT, + "linear1": gguf.MODEL_TENSOR.FFN_UP, + "linear2": gguf.MODEL_TENSOR.FFN_DOWN, + } + + def set_vocab(self): + # this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do + # unigram segmentation, so use the UGM tokenizer instead + from sentencepiece import sentencepiece_model_pb2 as model + + proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute] + proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read()) + assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer" + + tokens, scores, toktypes = self._create_vocab_sentencepiece() + + # the last rows of the embedding table are not sentencepiece pieces + extra = self._extra_tokens() + for i, name in enumerate(extra): + tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") + toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL + scores[len(tokens) - len(extra) + i] = -1000.0 + + self.gguf_writer.add_tokenizer_model("t5") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix) + self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces) + if proto.normalizer_spec.precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap) + self.gguf_writer.add_add_bos_token(False) + self.gguf_writer.add_add_eos_token(False) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if not name.startswith("flow_lm."): + return # mimi and the flow net go to the mmproj + + if name == "flow_lm.conditioner.embed.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch)) + return + + if name.startswith("flow_lm.out_norm."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.transformer.layers."): + assert bid is not None + key_with_suffix = name.split(f"layers.{bid}.", 1)[1] + key, suffix = key_with_suffix.rsplit(".", 1) + + if key == "self_attn.in_proj": + q, k, v = data_torch.chunk(3, dim=0) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v) + return + + tensor = self._LAYER_TENSOR_MAP.get(key) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch) + return + + return + + def _extra_tokens(self) -> list[str]: + # the conditioner's padding row, then the learned vectors appended by _embd_table(). + # bos_before_voice only exists when the pack sets insert_bos_before_voice + names = ["<|pad|>"] + if "flow_lm.bos_before_voice" in self.model_tensors: + names.append("<|bos_before_voice|>") + names.append("<|audio_bos|>") + return names + + def _embd_table(self, embed: Tensor) -> Tensor: + rows = [embed] + if "flow_lm.bos_before_voice" in self.model_tensors: + rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype)) + + # bos_emb is a latent, it only enters the backbone through input_linear + bos_emb = self.model_tensors["flow_lm.bos_emb"]() + input_linear = self.model_tensors["flow_lm.input_linear.weight"]() + audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1) + rows.append(audio_bos.to(embed.dtype)) + + return torch.cat(rows, dim=0) + + +@ModelBase.register("PocketTTSModel") +class PocketTTSMmprojModel(MmprojModel): + has_audio_encoder = True + has_vision_encoder = False + + _MIMI_TFM_MAP = { + "norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM), + "norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM), + "self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT), + "linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP), + "linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN), + "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE), + "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), + } + _MIMI_TFM_QKV = ( + (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V), + (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V), + ) + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + assert self.hparams_audio is not None + + # voice-prompt encoder: mimi encoder + speaker_proj + self.gguf_writer.add_clip_has_audio_encoder(True) + # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC) + self.gguf_writer.add_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + # mimi convolves the waveform directly, it is passed around as a 1-row "mel" + self.gguf_writer.add_audio_num_mel_bins(1) + + # generation: flow-matching decoder + mimi decoder + # the SEANet and flow net hparams are constant across the family, clip.cpp holds them + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN) + self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + + self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del name, bid, n_dims + # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path + if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"): + return gguf.GGMLQuantizationType.F16 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # the block index of the mimi transformers is parsed here, not by the base class + T = gguf.MODEL_TENSOR + + if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"): + return # folded into the backbone embedding table + if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."): + return # backbone + + if name == "flow_lm.speaker_proj_weight": + yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch) + return + if name == "flow_lm.input_linear.weight": + yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch) + return + if name == "flow_lm.emb_mean": + yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch) + return + if name == "flow_lm.emb_std": + yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch) + return + if name.startswith("flow_lm.out_eos."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.flow_net."): + yield from self._flow_net_tensor(name, data_torch) + return + + if name == "mimi.downsample.conv.conv.weight": + yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch) + return + if name == "mimi.upsample.convtr.convtr.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch) + return + if name == "mimi.quantizer.output_proj.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1)) + return + + if "_transformer.transformer.layers." in name: + yield from self._mimi_tfm_tensor(name, data_torch) + return + + if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."): + yield from self._seanet_tensor(name, data_torch) + return + + return + + def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + key = name.split("flow_lm.flow_net.", 1)[1] + suffix = "." + key.rsplit(".", 1)[1] + + simple = { + "input_proj": T.A_GEN_FLOW_INPUT_PROJ, + "cond_embed": T.A_GEN_FLOW_COND_EMBD, + "final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ, + "final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA, + } + tensor = simple.get(key.rsplit(".", 1)[0]) + if tensor is not None: + yield (self.format_tensor_name(tensor, suffix=suffix), data_torch) + return + + if key.startswith("time_embed."): + bid = int(key.split(".")[1]) + rest = key.split(f"time_embed.{bid}.", 1)[1] + time_map = { + "freqs": (T.A_GEN_FLOW_TIME_FREQS, ""), + "mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix), + "mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix), + "mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""), + } + entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0]) + if entry is not None: + yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch) + return + + if key.startswith("res_blocks."): + bid = int(key.split(".")[1]) + rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0] + blk_map = { + "in_ln": T.A_GEN_FLOW_BLK_NORM, + "mlp.0": T.A_GEN_FLOW_BLK_UP, + "mlp.2": T.A_GEN_FLOW_BLK_DOWN, + "adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA, + } + tensor = blk_map.get(rest) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + return + + def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + is_decoder = name.startswith("mimi.decoder_transformer.") + bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0]) + key_with_suffix = name.split(f".layers.{bid}.", 1)[1] + + if key_with_suffix == "self_attn.in_proj.weight": + q, k, v = data_torch.chunk(3, dim=0) + names = self._MIMI_TFM_QKV[1 if is_decoder else 0] + for tensor, part in zip(names, (q, k, v)): + yield (self.format_tensor_name(tensor, bid), part) + return + + key, suffix = key_with_suffix.rsplit(".", 1) + entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix) + if entry is None: + return + tensor = entry[1 if is_decoder else 0] + suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + + def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + is_decoder = name.startswith("mimi.decoder.") + idx = int(name.split(".model.", 1)[1].split(".")[0]) + suffix = "." + name.rsplit(".", 1)[1] + + conv_in, conv_out, res1, res2, scale = ( + (T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1, + T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV) + if is_decoder else + (T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1, + T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV) + ) + + if idx == 0: + yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch) + return + if idx == 3 * _N_SEANET_STAGES + 2: + yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch) + return + + for stage in range(_N_SEANET_STAGES): + res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) + scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage) + if idx == scale_idx: + yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch) + return + if idx == res_idx: + # block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU) + inner = int(name.split(".block.", 1)[1].split(".")[0]) + tensor = res1 if inner == 1 else res2 + yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch) + return diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8d0cad59b2cd..a197cacd3fa5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -407,6 +407,8 @@ class Projector: class ClipGenAudio: PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models + # name of the weight variant, for settings that are not in the checkpoint + MODEL_VARIANT = "clip.gen.audio.model_variant" EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" @@ -581,6 +583,7 @@ class MODEL_ARCH(IntEnum): MELLUM = auto() NANBEIGE = auto() QWEN3TTS = auto() + POCKETTTS = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1040,6 +1043,38 @@ class MODEL_TENSOR(IntEnum): A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM + # pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path) + A_ENC_SEANET_CONV_IN = auto() + A_ENC_SEANET_CONV_OUT = auto() + A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv + A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv + A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv + A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output + A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd + A_GEN_FLOW_INPUT_PROJ = auto() + A_GEN_FLOW_COND_EMBD = auto() + A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies + A_GEN_FLOW_TIME_UP = auto() + A_GEN_FLOW_TIME_DOWN = auto() + A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha + A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln + A_GEN_FLOW_BLK_UP = auto() + A_GEN_FLOW_BLK_DOWN = auto() + A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate + A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale + A_GEN_FLOW_FINAL_PROJ = auto() + A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state + A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd + A_GEN_EMB_MEAN = auto() # latent denormalization stats + A_GEN_EMB_STD = auto() + A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim + A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr + A_GEN_WAV_SEANET_CONV_IN = auto() + A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM + A_GEN_WAV_SEANET_RES_CONV1 = auto() + A_GEN_WAV_SEANET_RES_CONV2 = auto() + A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr A_MMPROJ = auto() A_MMPROJ_FC = auto() A_MM_NORM_PRE = auto() @@ -1255,6 +1290,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", MODEL_ARCH.QWEN3TTS: "qwen3tts", + MODEL_ARCH.POCKETTTS: "pockettts", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1709,6 +1745,37 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2", MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake", MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv", + MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in", + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv", + MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1", + MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2", + MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj", + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj", + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd", + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs", + MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj", + MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos", + MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear", + MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean", + MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std", + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out", + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv", MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", @@ -2020,6 +2087,37 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2, MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE, MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV, + MODEL_TENSOR.A_ENC_SEANET_CONV_IN, + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2, + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV, + MODEL_TENSOR.A_ENC_ATTN_SCALE, + MODEL_TENSOR.A_ENC_FFN_SCALE_LS, + MODEL_TENSOR.A_ENC_SPEAKER_PROJ, + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ, + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD, + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS, + MODEL_TENSOR.A_GEN_FLOW_TIME_UP, + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN, + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_UP, + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN, + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ, + MODEL_TENSOR.A_GEN_OUT_EOS, + MODEL_TENSOR.A_GEN_INPUT_LINEAR, + MODEL_TENSOR.A_GEN_EMB_MEAN, + MODEL_TENSOR.A_GEN_EMB_STD, + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT, + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV, MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, MODEL_TENSOR.A_ENC_CONV_NORM_VAR, MODEL_TENSOR.A_ENC_MEL_FILTERS, @@ -4903,6 +5001,18 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.POCKETTTS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], } # tensors that will not be serialized @@ -5179,6 +5289,8 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder + POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 81ae07c11a60..05f86396dc0a 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1453,6 +1453,9 @@ def add_gen_audio_head_count_kv(self, value: int) -> None: def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value) + def add_gen_audio_model_variant(self, value: str) -> None: + self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value) + def add_xielu_alpha_p(self, values: Sequence[float]): self.add_array(Keys.xIELU.ALPHA_P, values) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 73fb8b981382..8ed9391d7c74 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -147,6 +147,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_QWEN3TTS, "qwen3tts" }, + { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index 51dfd288ced2..18d9de186f75 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -152,6 +152,7 @@ enum llm_arch { LLM_ARCH_DFLASH, LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, + LLM_ARCH_POCKETTTS, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index cc1917b6da77..0e27cb41713b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -116,6 +116,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen3vlmoe(params); case LLM_ARCH_QWEN3TTS: return new llama_model_qwen3tts(params); + case LLM_ARCH_POCKETTTS: + return new llama_model_pockettts(params); case LLM_ARCH_PHI2: return new llama_model_phi2(params); case LLM_ARCH_PHI3: @@ -2637,6 +2639,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MAINCODER: case LLM_ARCH_GLM_DSA: case LLM_ARCH_NANBEIGE: + case LLM_ARCH_POCKETTTS: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 diff --git a/src/models/models.h b/src/models/models.h index 9230345208a9..ddb9ae2f1210 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -713,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base { }; +struct llama_model_pockettts : public llama_model_base { + llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_codeshell : public llama_model_base { llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/pockettts.cpp b/src/models/pockettts.cpp new file mode 100644 index 000000000000..1b3bb6c648af --- /dev/null +++ b/src/models/pockettts.cpp @@ -0,0 +1,146 @@ +#include "models.h" + +// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model. +// it has no lm_head, the audio latents are produced by the flow net inside the mmproj + +void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + + switch (hparams.n_layer()) { + case 6: type = LLM_TYPE_109M; break; + case 24: type = LLM_TYPE_335M; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_pockettts::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0); + // no output head, the logits are unused; reuse the embedding table so a sampler can still run + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0); + + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } +} + +std::unique_ptr llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + cur = build_norm(inpL, + model.layers[il].attn_norm, + model.layers[il].attn_norm_b, + LLM_NORM, il); + cb(cur, "attn_norm", il); + + // self-attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + // FF + { + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, + model.layers[il].ffn_norm_b, + LLM_NORM, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_GELU, LLM_FFN_SEQ, il); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = build_norm(inpL, + model.output_norm, + model.output_norm_b, + LLM_NORM, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index fe22cb12543e..3f4a6c670dd7 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -57,6 +57,9 @@ add_library(mtmd models/mimo-audio.cpp models/qwen3tts-spkenc.cpp models/qwen3tts-gen.cpp + models/pockettts-seanet.cpp + models/pockettts-spkenc.cpp + models/pockettts-gen.cpp models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index 3cddd085ec6c..ac43e1b81b1c 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i ### Checklist for porting new audio generation models to mtmd -1. Establish a list of reusable and missing components from the current mtmd implementation. -2. For GGUF conversion: +1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments + - Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged +2. Establish a list of reusable and missing components from the current mtmd implementation. +3. For GGUF conversion: - Backbone model should be converted to a normal text model (loadable via `libllama`) - If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`) - If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`) @@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i - For tensor naming: - Prefixed with `a.*` for tensors used by speaker encoder pipeline - Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation) -3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: + - For GGUF metadata: + - Reuse as many existing keys as possible + - In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams` + - If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary + - Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them +4. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: - 10-20% changes is to add new backbone (text) model and conversion - 60% changes inside `mtmd-helper-gen.cpp` - 10% changes inside `libmtmd` and `clip.cpp` systems - The rest downstream code (CLI, server) should have no changes at all -4. Update usage documentation in `tools/tts/README.md` +5. Update usage documentation in `tools/tts/README.md` IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**. diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index bf73222f9def..b2c8b4021321 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -92,7 +92,9 @@ #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size // audio generation (gen-audio)-specific #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities -#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor" +// name of the weight variant, for settings that are not in the checkpoint +#define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant" +#define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // // tensor name constants @@ -246,6 +248,38 @@ #define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s" #define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s" +// pocket-tts +#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s" +#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s" +#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s" +#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s" +#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s" +#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s" +#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s" +#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s" +#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s" +#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs" +#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s" +#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s" +#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm" +#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s" +#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s" +#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s" +#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s" +#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s" +#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s" +#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s" +#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s" +#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean" +#define TN_A_GEN_EMB_STD "a.gen.emb_std" +#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s" +#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s" +#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s" +#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s" +#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s" + // cogvlm #define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s" #define TN_MM_H_TO_4H "mm.up.%s" @@ -455,6 +489,8 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_POCKETTTS_SPKENC, + PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_UNKNOWN, }; @@ -515,6 +551,8 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, + { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 761aabf64e83..ad25c008e738 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -141,6 +141,20 @@ struct clip_hparams { int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) + // threshold for the "out_eos_score" graph output + float gen_eos_threshold = 0.0f; + + // name of the weight variant, some pipelines tune themselves on it + std::string gen_model_variant; + + // pocket-tts + static constexpr int32_t pockettts_max_spk_seconds = 30; + int32_t seanet_n_stage = 0; + std::vector seanet_ratios; // encoder order (reversed compared to the config) + int32_t mimi_downsample = 0; // encoder frame rate / model frame rate + int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames + int32_t flow_n_step = 1; // lsd_decode steps + // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; int32_t wav_tfm_n_embd = 0; @@ -402,6 +416,63 @@ struct qf_block { std::vector qf_proj_layers; }; +// pocket-tts SEANet stack, used in both directions: +// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out +// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out +struct clip_seanet { + // one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input + struct stage { + ggml_tensor * res_conv1_w = nullptr; + ggml_tensor * res_conv1_b = nullptr; + ggml_tensor * res_conv2_w = nullptr; + ggml_tensor * res_conv2_b = nullptr; + ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder) + ggml_tensor * scale_conv_b = nullptr; + }; + + ggml_tensor * conv_in_w = nullptr; + ggml_tensor * conv_in_b = nullptr; + ggml_tensor * conv_out_w = nullptr; + ggml_tensor * conv_out_b = nullptr; + std::vector stages; +}; + +// pocket-tts flow-matching decoder (SimpleMLPAdaLN) +struct clip_flow_net { + // AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual + struct block { + ggml_tensor * norm_w = nullptr; + ggml_tensor * norm_b = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * ada_w = nullptr; // -> shift, scale, gate + ggml_tensor * ada_b = nullptr; + }; + + // timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm + struct time_embd { + ggml_tensor * freqs = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * norm = nullptr; // RMSNorm alpha + }; + + ggml_tensor * input_proj_w = nullptr; + ggml_tensor * input_proj_b = nullptr; + ggml_tensor * cond_embd_w = nullptr; + ggml_tensor * cond_embd_b = nullptr; + ggml_tensor * final_ada_w = nullptr; // -> shift, scale + ggml_tensor * final_ada_b = nullptr; + ggml_tensor * final_proj_w = nullptr; + ggml_tensor * final_proj_b = nullptr; + std::vector time; + std::vector blocks; +}; + // qwen3tts code2wav: RVQ codes -> raw PCM struct clip_code2wav { // "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it @@ -699,6 +770,24 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path) + clip_seanet seanet; + + // pocket-tts: voice latent -> backbone embd (speaker path) + ggml_tensor * spk_proj_w = nullptr; + ggml_tensor * downsample_w = nullptr; + + // pocket-tts: flow-matching decoder, backbone hidden state -> next latent + clip_flow_net flow; + ggml_tensor * gen_out_eos_w = nullptr; + ggml_tensor * gen_out_eos_b = nullptr; + ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd + ggml_tensor * gen_emb_mean = nullptr; + ggml_tensor * gen_emb_std = nullptr; + ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim + ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate + std::vector gen_tfm_layers; // mimi decoder_transformer + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 1e53eddf8179..2fb2b5041dcd 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -174,6 +174,10 @@ struct clip_ctx { bool support_batch = false; + // for audio gen, reseeded only when the caller asks for another seed + std::mt19937 rng{std::random_device{}()}; + uint32_t rng_seed = UINT32_MAX; + clip_ctx(clip_context_params & ctx_params) { flash_attn_type = ctx_params.flash_attn_type; no_alloc = ctx_params.no_alloc; @@ -1059,6 +1063,25 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; + const int n_step = ctx->model.hparams.flow_n_step; + const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; + GGML_ASSERT(n_step > 0); + GGML_ASSERT(n_latent > 0); + // "inp_feats" takes the caller's buffer as-is, the graph must consume all of it + if (params && params->feats) { + GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0); + GGML_ASSERT(params->feats->size() >= (size_t) n_latent); + } + const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; + builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; @@ -1282,6 +1305,7 @@ struct clip_model_loader { // these are unused, but still need to be set to avoid issues hparams.image_size = 0; hparams.patch_size = 1; + get_string(KEY_GEN_AUDIO_VARIANT, hparams.gen_model_variant, false); } else { GGML_ASSERT(false && "unknown modality"); @@ -1421,7 +1445,7 @@ struct clip_model_loader { } break; case PROJECTOR_TYPE_PARAKEET: { - get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor); + get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor); GGML_ASSERT(hparams.subsampling_factor == 8 && "subsampling_factor must match the conv strides in clip_graph_parakeet::build()"); get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); @@ -1745,6 +1769,22 @@ struct clip_model_loader { // matches the reference decoder's sliding_window (speech_tokenizer/config.json) hparams.wav_tfm_swa = 72; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // mimi front-end takes the raw waveform, no mel + hparams.audio_sample_rate = 24000; + // seanet ratios are [6,5,4] in the config, the encoder reverses them + hparams.seanet_ratios = { 4, 5, 6 }; + hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size(); + hparams.mimi_downsample = 16; + // matches the reference transformer's "context" + hparams.mimi_tfm_context = 250; + hparams.rope_theta = 10000.0f; + // flow_lm defaults, see pocket_tts/default_parameters.py + hparams.flow_n_step = 1; + hparams.gen_eos_threshold = -4.0f; + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -1947,7 +1987,9 @@ struct clip_model_loader { // GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT // checks below do not apply to it. - const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA; + // pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform + const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC; // Validate audio hparams loaded from GGUF metadata if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) { @@ -2020,6 +2062,31 @@ struct clip_model_loader { return cur; }; + // pocket-tts: the encoder and the decoder share the same layout, only the prefix differs + auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) { + const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN; + const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT; + const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1; + const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2; + const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV; + + seanet.conv_in_w = get_tensor(string_format(conv_in, "weight")); + seanet.conv_in_b = get_tensor(string_format(conv_in, "bias")); + seanet.conv_out_w = get_tensor(string_format(conv_out, "weight")); + seanet.conv_out_b = get_tensor(string_format(conv_out, "bias")); + + seanet.stages.resize(hparams.seanet_n_stage); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + auto & stage = seanet.stages[i]; + stage.res_conv1_w = get_tensor(string_format(res1, i, "weight")); + stage.res_conv1_b = get_tensor(string_format(res1, i, "bias")); + stage.res_conv2_w = get_tensor(string_format(res2, i, "weight")); + stage.res_conv2_b = get_tensor(string_format(res2, i, "bias")); + stage.scale_conv_w = get_tensor(string_format(scale, i, "weight")); + stage.scale_conv_b = get_tensor(string_format(scale, i, "bias")); + } + }; + auto get_vector = [&](const std::string & name) { std::vector result; auto it = tensor_offset.find(name); @@ -2081,7 +2148,8 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && - model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC); + model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2755,6 +2823,81 @@ struct clip_model_loader { model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight")); model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias")); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + load_seanet(model.seanet, false); + model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight")); + model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight")); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + auto & flow = model.flow; + flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight")); + flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias")); + flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight")); + flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias")); + flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight")); + flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias")); + flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight")); + flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias")); + + flow.time.resize(2); + for (size_t i = 0; i < flow.time.size(); i++) { + auto & t = flow.time[i]; + t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i)); + t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight")); + t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias")); + t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight")); + t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias")); + t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i)); + } + + // one AdaLN block per flow depth, the count is only known from the tensors + for (int il = 0; ; il++) { + ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false); + if (probe == nullptr) { + break; + } + clip_flow_net::block blk; + blk.norm_w = probe; + blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias")); + blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight")); + blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias")); + blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight")); + blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias")); + blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight")); + blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias")); + flow.blocks.push_back(blk); + } + + model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight")); + model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias")); + model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight")); + model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN); + model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD); + + // mimi decoder + model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight")); + model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight")); + load_seanet(model.seanet, true); + model.gen_tfm_layers.resize(hparams.n_layer); + for (int il = 0; il < hparams.n_layer; il++) { + auto & layer = model.gen_tfm_layers[il]; + const char * p = "a.gen.wav.tfm"; + layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight")); + layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias")); + layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight")); + layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight")); + layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight")); + layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight")); + layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight")); + layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight")); + layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias")); + layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight")); + layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight")); + layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight")); + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // code_predictor @@ -4060,6 +4203,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // one hidden-state vector fed back to the talker per call n_patches = 1; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + // one conditioning row per 12.5Hz frame + const int hop = ctx->model.hparams.mimi_downsample * 120; + n_patches = img->nx() / hop; + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller + n_patches = 1; + } break; case PROJECTOR_TYPE_GRANITE4_VISION: { // Per-tile output token count: each projector block outputs @@ -4101,6 +4255,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } +// persisted state slots of the gen-audio decoder, per pipeline +static std::vector list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) { + switch (model.proj_type) { + case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model); + case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model); + default: return {}; + } +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4116,6 +4279,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { clip_model_loader::warmup(*ctx, *params->imgs); } + if (params->seed != ctx->rng_seed) { + ctx->rng_seed = params->seed; + ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed); + } + // build the inference graph ggml_backend_sched_reset(ctx->sched.get()); ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build(); @@ -4160,6 +4328,50 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur)); }; + // upload the decoder state from the previous call, or zero-fill on a cold start + auto set_gen_state_in = [&]() { + size_t offset = 0; + for (const auto & slot : list_gen_state_slots(hparams, model)) { + ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); + const size_t nb = ggml_nbytes(t); + if (params->state_in && params->state_in->size() >= offset + nb) { + ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); + } else { + std::vector zeros(nb, 0); + ggml_backend_tensor_set(t, zeros.data(), 0, nb); + } + offset += nb; + } + }; + + // rope positions and attention mask of the mimi transformers (pocket-tts). + // the mask is causal with a sliding window, see _build_attention_mask() in the reference + auto set_pockettts_tfm_inputs = [&]() { + const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos")); + GGML_ASSERT(n_pos > 0); + std::vector positions((size_t) n_pos); + for (int64_t i = 0; i < n_pos; i++) { + positions[(size_t) i] = (int32_t) i; + } + set_input_i32("inp_pos", positions); + + // the preprocessor truncates the waveform to keep this mask bounded + const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120; + GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask"); + + const int64_t context = hparams.mimi_tfm_context; + std::vector mask((size_t) n_pos * n_pos, -INFINITY); + for (int64_t q = 0; q < n_pos; q++) { + for (int64_t k = 0; k < n_pos; k++) { + const int64_t delta = q - k; + if (delta >= 0 && delta < context) { + mask[(size_t) q * n_pos + k] = 0.0f; + } + } + } + set_input_f32("kq_mask", mask); + }; + // set input pixel values if (!imgs.is_audio) { size_t nelem = 0; @@ -4203,8 +4415,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) { - // audio input, code2wav is not here: its only input is "inp_codes", set in the switch below + } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) { + // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below GGML_ASSERT(imgs.entries.size() == 1); const auto & mel_inp = imgs.entries[0]; @@ -4737,6 +4949,30 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_i32("patches", patches); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + set_pockettts_tfm_inputs(); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { + GGML_ASSERT(params->feats != nullptr); + set_input_f32("inp_feats", *params->feats); + // positions and mask are derived in-graph from the persisted counter + set_gen_state_in(); + } else { + // flow matching starts from gaussian noise, std = sqrt(temp) + ggml_tensor * t = get_inp_tensor("inp_noise"); + // Config.default_temperature, for a caller that does not set one + const float temp = params->temp > 0.0f ? params->temp : 0.7f; + std::normal_distribution dist(0.0f, std::sqrt(temp)); + std::vector noise(ggml_nelements(t)); + for (auto & v : noise) { + v = dist(ctx->rng); + } + set_input_f32("inp_noise", noise); + } + } break; case PROJECTOR_TYPE_GEMMA4V: case PROJECTOR_TYPE_GEMMA4UV: { @@ -4861,20 +5097,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } set_input_i32("inp_codes", codes); - - // upload the state from the previous call, or zero-fill on a cold start - size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { - ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); - const size_t nb = ggml_nbytes(t); - if (params->state_in && params->state_in->size() >= offset + nb) { - ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); - } else { - std::vector zeros(nb, 0); - ggml_backend_tensor_set(t, zeros.data(), 0, nb); - } - offset += nb; - } + set_gen_state_in(); } else { // code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it const int64_t vocab0 = model.gen_code_out_embd_w->ne[1]; @@ -4886,11 +5109,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_i32("inp_code0", code0); // one uniform(0,1) draw per codebook, used by do_sampling() - static std::mt19937 rng{ std::random_device{}() }; std::uniform_real_distribution dist(0.0f, 1.0f); const int64_t n_acoustic = model.gen_code_head_w->ne[2]; for (int64_t g = 0; g < n_acoustic; g++) { - std::vector r = { dist(rng) }; + std::vector r = { dist(ctx->rng) }; set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); } } @@ -5343,14 +5565,31 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // for audio gen models // + // optional outputs: a pipeline yields codes or feats, and not all have an eos head if (params->out_codes != nullptr) { ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); - if (codes == nullptr) { - GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor"); + if (codes != nullptr) { + auto & out_codes = *params->out_codes; + out_codes.resize(ggml_nelements(codes)); + ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); + } + } + if (params->out_feats != nullptr) { + ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats"); + if (feats != nullptr) { + auto & out_feats = *params->out_feats; + out_feats.resize(ggml_nelements(feats)); + ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats)); + } + } + if (params->out_is_eos != nullptr) { + ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score"); + if (eos != nullptr) { + GGML_ASSERT(ggml_nelements(eos) == 1); + float score = 0.0f; + ggml_backend_tensor_get(eos, &score, 0, sizeof(float)); + *params->out_is_eos = score > hparams.gen_eos_threshold; } - auto & out_codes = *params->out_codes; - out_codes.resize(ggml_nelements(codes)); - ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } if (params->out_audio != nullptr) { ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio"); @@ -5362,9 +5601,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); // drop the tail audio that comes from the code-0 rear padding - const int64_t n_codes = model.gen_code_head_w->ne[2] + 1; + const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0; const int64_t n_frames_w = hparams.wav_tfm_swa; - const int64_t n_frames = (int64_t) params->codes->size() / n_codes; + const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w; if (n_frames < n_frames_w) { const size_t hop = out_audio.size() / n_frames_w; out_audio.resize((size_t) n_frames * hop); @@ -5373,12 +5612,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { if (params->state_out != nullptr) { auto & state_out = *params->state_out; size_t total = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float); } state_out.resize(total); size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str()); if (t == nullptr) { GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str()); @@ -5526,6 +5765,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + return ctx->model.spk_proj_w->ne[1]; + case PROJECTOR_TYPE_POCKETTTS_GEN: + return ctx->model.gen_input_lin_w->ne[1]; case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 7f706d976eb3..a5b713775234 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -104,9 +104,14 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; // this frame's 16 sampled codes + std::vector * out_feats = nullptr; // continuous counterpart of out_codes + uint32_t seed = UINT32_MAX; // UINT32_MAX for random + float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders + bool * out_is_eos = nullptr; // GEN_WAV const std::vector * codes = nullptr; // this frame's 16 RVQ codes + const std::vector * feats = nullptr; // continuous counterpart of codes std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 519c0d019add..ed8c1ea51870 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -318,6 +318,59 @@ struct clip_graph_qwen3tts_gen : clip_graph { }; }; +// +// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder. +// stateless unless state_in is populated: convs then pad instead of carrying left-context. +// +struct clip_graph_pockettts_seanet : clip_graph { + clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {} + ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); } + + // per-call streaming state, keyed by slot name (see list_pockettts_state_slots) + std::map state_in; + mutable std::vector> state_out; + + ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate = false, const std::string & state_name = "") const; + ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name = "") const; + ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix = "") const; + + // x: [T, C] -> [T / hop, dim] + ggml_tensor * encode(ggml_tensor * x) const; + // x: [T, dim] -> [T * hop, 1], streams when state_in is populated + ggml_tensor * decode(ggml_tensor * x) const; +}; + +// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows +struct clip_graph_pockettts_spkenc : clip_graph { + clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + + ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const; +}; + +// +// pocket-tts generation: +// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call +// GEN_WAV = mimi decoder, a window of latents -> PCM +// +struct clip_graph_pockettts_gen : clip_graph { + clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames) + : clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {} + ggml_cgraph * build() override; + + clip_gen_process_type gen_process; + int n_step; // lsd_decode steps, fixed at graph-build time + int n_frames; // GEN_WAV only: number of latents to decode + + // AdaLN modulation: x * (1 + scale) + shift + ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const; + ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const; + ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const; +}; + // one persisted state buffer used by code2wav, see qwen3tts-gen.cpp struct c2w_state_slot { std::string name; @@ -326,6 +379,9 @@ struct c2w_state_slot { }; std::vector list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model); +// same, for the streaming mimi decoder (pocket-tts GEN_WAV) +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model); + struct clip_graph_kimik25 : clip_graph { clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp new file mode 100644 index 000000000000..3fd613e5f7fa --- /dev/null +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -0,0 +1,291 @@ +#include "models.h" + +#include + +// pocket-tts generation stages +// +// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score +// GEN_WAV : a window of latents -> PCM, through the mimi decoder +// +// there is no codebook anywhere, "codes" in the mtmd API are continuous features here + +ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const { + ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f)); + return ggml_add(ctx0, cur, shift); +} + +// see TimestepEmbedder in the reference +ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const { + // t is a graph-build constant, so the cos/sin table can be folded into a scaled copy + ggml_tensor * args = ggml_scale(ctx0, te.freqs, t); + ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0); + + ggml_tensor * cur = build_mm(te.up_w, emb); + cur = ggml_add(ctx0, cur, te.up_b); + cur = ggml_silu(ctx0, cur); + cur = build_mm(te.down_w, cur); + cur = ggml_add(ctx0, cur, te.down_b); + + // this "RMSNorm" divides by the unbiased variance, not the mean square + // it also rescales the input, not the centered value, see _rms_norm() in mlp.py + { + const int64_t n = cur->ne[0]; + ggml_tensor * mean = ggml_mean(ctx0, cur); + ggml_tensor * dev = ggml_sub(ctx0, cur, mean); + ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev)); + var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f); + cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var)); + cur = ggml_mul(ctx0, cur, te.norm); + } + + return cur; +} + +// one velocity evaluation: v(cond, s, t, x) +ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const { + const auto & flow = model.flow; + + ggml_tensor * cur = build_mm(flow.input_proj_w, x); + cur = ggml_add(ctx0, cur, flow.input_proj_b); + + // the two time conditions are averaged, then added to the projected backbone state + ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t)); + ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size()); + + ggml_tensor * c = build_mm(flow.cond_embd_w, cond); + c = ggml_add(ctx0, c, flow.cond_embd_b); + + ggml_tensor * y = ggml_add(ctx0, ts, c); + cb(y, "flow_cond", -1); + + const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0]; + + for (size_t il = 0; il < flow.blocks.size(); il++) { + const auto & blk = flow.blocks[il]; + + ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, blk.ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]); + + ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il); + h = modulate(h, shift, scale); + h = build_mm(blk.up_w, h); + h = ggml_add(ctx0, h, blk.up_b); + h = ggml_silu(ctx0, h); + h = build_mm(blk.down_w, h); + h = ggml_add(ctx0, h, blk.down_b); + + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h)); + cb(cur, "flow_blk", (int) il); + } + + // final layer: the norm has no weights, only the AdaLN modulation + ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, flow.final_ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + + cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1); + cur = modulate(cur, shift, scale); + cur = build_mm(flow.final_proj_w, cur); + cur = ggml_add(ctx0, cur, flow.final_proj_b); + + return cur; +} + +// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context +// and the transposed-conv overlap tails +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector slots; + if (model.gen_upsample_w == nullptr) { + return slots; // not a pocket-tts decoder + } + const auto & seanet = model.seanet; + + // the slots below are sized from these + GGML_ASSERT(!model.gen_tfm_layers.empty()); + GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage); + GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage); + GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0); + + slots.push_back({"tfm_pos", 1, 1}); + + const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) { + slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix}); + slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix}); + } + + // upsample is depthwise, its output channel count is the input one + slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]}); + + slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]}); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]}); + slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]}); + } + slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]}); + + return slots; +} + +ggml_cgraph * clip_graph_pockettts_gen::build() { + if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) { + // the backbone hidden state arrives as the single batch entry + ggml_tensor * h_state = build_inp_raw(1); + h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1); + + // end-of-speech probe, thresholded on the host side + ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state); + eos = ggml_add(ctx0, eos, model.gen_out_eos_b); + ggml_set_name(eos, "out_eos_score"); + ggml_set_output(eos); + ggml_build_forward_expand(gf, eos); + + const int64_t n_latent = model.gen_input_lin_w->ne[0]; + + ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1); + ggml_set_name(noise, "inp_noise"); + ggml_set_input(noise); + + // lsd_decode: integrate the velocity field from the noise sample + ggml_tensor * cur = noise; + for (int i = 0; i < n_step; i++) { + const float s = (float) i / (float) n_step; + const float t = (float) (i + 1) / (float) n_step; + ggml_tensor * v = flow_forward(h_state, cur, s, t); + cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step)); + } + cb(cur, "flow_latent", -1); + + ggml_set_name(cur, "out_feats"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + // the same latent, projected into the backbone's input space for the next step + ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur); + cb(embd, "gen_embd", -1); + ggml_build_forward_expand(gf, embd); + + return gf; + } + + // GEN_WAV: [32, n_frames] latents -> PCM + ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, + model.gen_input_lin_w->ne[0], n_frames); + ggml_set_name(feats, "inp_feats"); + ggml_set_input(feats); + + // denormalize, then the DummyQuantizer up-projection + ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean); + cur = build_mm(model.gen_quant_out_w, cur); + cb(cur, "quant_out", -1); + + clip_graph_pockettts_seanet seanet(*this); + for (const auto & slot : list_pockettts_state_slots(hparams, model)) { + ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_set_name(t, ("state_in_" + slot.name).c_str()); + ggml_set_input(t); + seanet.state_in[slot.name] = t; + } + + // model frame rate -> encoder frame rate, depthwise transposed conv + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up"); + cb(cur, "mimi_upsample", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // positions continue across calls, the counter lives in the state + const int64_t n_pos = cur->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + const int64_t n_kv = prefix + n_pos; + + ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1); + ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base), + GGML_TYPE_I32); + seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)}); + + // banded causal mask over [cached prefix | this chunk] + // the last factor masks out cache rows that hold no real frame yet + ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1); + ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos); + ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k); + + ggml_tensor * keep = ggml_mul(ctx0, + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0 + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context + keep = ggml_mul(ctx0, keep, + ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix))); + ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1); + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.gen_tfm_layers[il]; + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + // prepend the cached window, then keep this chunk's tail for the next call + const std::string k_name = "tfm_k_" + std::to_string(il); + const std::string v_name = "tfm_v_" + std::to_string(il); + ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name), + ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1); + ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1); + seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix, + k_full->nb[1], (size_t) n_pos * k_full->nb[1]))}); + seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix, + v_full->nb[1], (size_t) n_pos * v_full->nb[1]))}); + + ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1); + ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1); + ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1); + + cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + } + cb(cur, "mimi_dec_tfm", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.decode(cur); + + for (const auto & s : seanet.state_out) { + ggml_set_name(s.second, ("state_out_" + s.first).c_str()); + ggml_set_output(s.second); + ggml_build_forward_expand(gf, s.second); + } + + // [n_samples, 1] -> [n_samples], clamped like the reference output + cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]); + cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f); + ggml_set_name(cur, "out_audio"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp new file mode 100644 index 000000000000..c47207f569bb --- /dev/null +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -0,0 +1,162 @@ +#include "models.h" + +// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py +// +// tensors are T-first here: [T, C] +// the convs are causal: left context comes from a state slot, or from padding on a cold start + +static int64_t div_ceil(int64_t a, int64_t b) { + return a / b + (a % b ? 1 : 0); +} + +// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC] +// the convs are causal, so the whole K - stride padding goes on the left +ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate, const std::string & state_name) const { + const int64_t k_size = (w->ne[0] - 1) * dilation + 1; + const int64_t p_total = k_size - stride; + + // trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py + const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride); + const int64_t ideal_len = n_frames * stride + k_size - p_total; + const int64_t p_extra = ideal_len - x->ne[0]; + + if (!state_name.empty() && p_total > 0) { + // streaming: the left context is the tail of the previous call + ggml_tensor * left = state_in.at(state_name); // [p_total, IC] + x = ggml_concat(ctx0, left, x, 0); + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1], + (size_t) (x->ne[0] - p_total) * x->nb[0]))}); + } else if (pad_replicate && p_total > 0) { + // the resamplers repeat the first frame instead of zero-padding + ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0); + ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1); + x = ggml_concat(ctx0, left, x, 0); + x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0); + } else { + x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0); + } + + ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation); + y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]); + if (b) { + y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return y; +} + +// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC] +// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped +ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name) const { + const int64_t K = w->ne[0]; + const int64_t T = x->ne[0]; + const int64_t p_total = K - stride; + const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1; + const int64_t OC = depthwise ? w->ne[2] : w->ne[1]; + const int64_t emit_len = T * stride; + + // one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride + ggml_tensor * col; + if (depthwise) { + // one group per channel: a batched matmul over the channels scales the kernel by each step + ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC] + ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC] + col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC] + col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T] + col = ggml_reshape_2d(ctx0, col, K * OC, T); + } else { + ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]); + w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC] + ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T] + col = ggml_mul_mat(ctx0, w2, xt); + } + ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC] + + ggml_tensor * out; + if (state_name.empty() || p_total == 0) { + out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0)); + } else { + // overlap-add the tail the previous call held back + ggml_tensor * prev = state_in.at(state_name); // [p_total, OC] + ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev); + if (emit_len > p_total) { + ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1], + (size_t) p_total * full->nb[0]); + out = ggml_concat(ctx0, head, rest, 0); + } else { + out = head; + } + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], + (size_t) emit_len * full->nb[0]))}); + } + + if (b) { + out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return out; +} + +ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix) const { + ggml_tensor * h = ggml_elu(ctx0, x); + h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix); + h = ggml_elu(ctx0, h); + // the second conv is pointwise, it needs no left context + h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1); + return ggml_add(ctx0, x, h); +} + +ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1); + cb(cur, "seanet_enc_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[i]; + + cur = res_unit(cur, stage, 1); + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1); + cb(cur, "seanet_enc_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1); + cb(cur, "seanet_enc_out", -1); + + return cur; +} + +ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + const bool stream = !state_in.empty(); + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false, + stream ? "dec_in" : ""); + cb(cur, "seanet_dec_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + // the decoder mirrors the encoder, so the ratios are walked backwards + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + const std::string id = std::to_string(i); + + cur = ggml_elu(ctx0, cur); + cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, + stream ? "dec_up_" + id : ""); + cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : ""); + cb(cur, "seanet_dec_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false, + stream ? "dec_out" : ""); + cb(cur, "seanet_dec_out", -1); + + return cur; +} diff --git a/tools/mtmd/models/pockettts-spkenc.cpp b/tools/mtmd/models/pockettts-spkenc.cpp new file mode 100644 index 000000000000..f802d90687d7 --- /dev/null +++ b/tools/mtmd/models/pockettts-spkenc.cpp @@ -0,0 +1,77 @@ +#include "models.h" + +// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame +// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight + +// pre-norm block with layer scale on both residual paths, see mimi_transformer.py +ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const { + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + const int64_t n_pos = cur->ne[1]; + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + + return cur; +} + +ggml_cgraph * clip_graph_pockettts_spkenc::build() { + // the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1] + ggml_tensor * inp_raw = build_inp_raw(1); + ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]); + + clip_graph_pockettts_seanet seanet(*this); + cur = seanet.encode(cur); + cb(cur, "mimi_enc", -1); + + // [T, 512] -> transformer works on [512, T] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + // the mimi transformer is causal with a sliding window, see _build_attention_mask() + ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]); + ggml_set_name(kq_mask, "kq_mask"); + ggml_set_input(kq_mask); + + for (int il = 0; il < n_layer; il++) { + cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il); + } + cb(cur, "mimi_enc_tfm", -1); + + // downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true); + cb(cur, "mimi_downsample", -1); + + // voice latent -> backbone embd + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = build_mm(model.spk_proj_w, cur); + cb(cur, "spk_proj", -1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index b6c95efa941e..84c77f4fad19 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -610,6 +610,10 @@ std::vector list_c2w_state_slots(const clip_hparams & hparams, c const auto & c2w = model.c2w; std::vector slots; + if (c2w.pre_conv_w == nullptr) { + return slots; // not a code2wav model, it keeps no state between calls + } + slots.push_back({"tfm_pos", 1, 1}); // prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward) diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ca4b64efa4a4..98a8c11ee915 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1423,3 +1423,41 @@ std::vector mtmd_audio_streaming_istft::flush() { return output; } + +// +// mtmd_audio_preprocessor_pockettts +// +// mimi takes the raw 24kHz waveform, there is no mel front-end +// the samples are handed over as a single-row "mel", to reuse the normal chunk path +// + +bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + // the encoder needs whole frames, see pad_for_conv1d() in the reference + const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120; + if (n_samples == 0 || frame_size <= 0) { + return false; + } + + // the mimi transformer mask is dense, so cost is quadratic in the reference length + const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate; + if ((int64_t) n_samples > max_samples) { + LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__, + (double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds); + n_samples = (size_t) max_samples; + } + + const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size; + const int64_t n_padded = n_frames * frame_size; + + mtmd_audio_mel out; + out.n_mel = 1; + out.n_len = n_padded; + out.n_len_org = (int64_t) n_samples; + out.data.assign((size_t) n_padded, 0.0f); + std::copy(samples, samples + n_samples, out.data.begin()); + + output.push_back(std::move(out)); + return true; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index b4d6f7259808..44ad098ae63a 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +// mimi convolves the waveform directly, so this only pads it to a whole number of frames +struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override {} + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; +}; + struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } void initialize() override; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 85671d1a331f..1c58d3ae1959 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -5,6 +5,8 @@ #include "../src/llama-ext.h" #include +#include +#include #include #include #include @@ -87,7 +89,8 @@ class mtmd_gen_audio_pipeline { virtual int32_t step_prompt(int32_t n_batch) = 0; // sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token, // those read what they need from h_state_in instead - virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0; + // set out_stop on end-of-speech, h_state_out must be null if no frame is generated + virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0; virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; protected: @@ -200,8 +203,10 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { prompt_pos = 0; pos = 0; - top_k = inp->top_k > 0 ? inp->top_k : 50; - top_p = inp->top_p > 0 ? inp->top_p : 1.0f; + const mtmd_gen_inp def = mtmd_gen_inp_default(mctx); + top_k = inp->top_k > 0 ? inp->top_k : def.top_k; + top_p = inp->top_p > 0 ? inp->top_p : def.top_p; + seed = inp->seed; out_type = inp->out_type; // the prompt above holds the whole text stream up to tts_eos, so every generated @@ -241,13 +246,26 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return n_prompt - prompt_pos; } - int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { - mtmd_gen_inp inp{}; + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + if (sampled == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n"); + return 1; + } + + // backbone signals end-of-speech with a token, no frame for this step + if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; inp.code0 = sampled - codec_0; inp.embd = const_cast(h_state_in); inp.top_k = top_k; inp.top_p = top_p; + inp.seed = seed; mtmd_gen_out out{}; if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n"); @@ -384,10 +402,11 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { if (codes_buf.empty()) { return true; } - mtmd_gen_inp inp{}; + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; inp.codes = codes_buf.data(); inp.n_codes = codes_buf.size(); + inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data(); inp.state_size = c2w_state.size(); mtmd_gen_out out{}; @@ -427,8 +446,9 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::unique_ptr prompt_batch; int n_prompt = 0; int prompt_pos = 0; - int32_t top_k = 50; - float top_p = 1.0f; + int32_t top_k = 50; + float top_p = 1.0f; + uint32_t seed = UINT32_MAX; std::vector codes_buf; std::vector c2w_state; std::vector audio_pcm; @@ -438,10 +458,547 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; +// settings that only live in the reference's per-pack yaml, not in the checkpoint +// the english packs share the same shapes and tokenizer, but disagree on these +// all three are 0 / false when the pack does not tune them, the model default is then used +struct pockettts_pack_settings { + float temp = 0.0f; + int frames_after_eos = 0; + bool pad_short_text = false; +}; + +static pockettts_pack_settings pockettts_pack(const char * variant) { + static const std::unordered_map packs = { + { "english", { 0.3f, 0, false } }, + { "english_2026-01", { 0.7f, 0, true } }, + { "english_2026-04", { 0.3f, 0, false } }, + { "french_24l", { 0.7f, 8, false } }, + }; + auto it = packs.find(variant ? variant : ""); + if (it == packs.end()) { + LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n", + variant ? variant : ""); + return {}; + } + return it->second; +} + +// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent +// the end-of-speech head also lives in the mmproj +class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + seq_id = 0; + pos = 0; + feats_buf.clear(); + dec_state.clear(); + audio_pcm.clear(); + h_state_buf.clear(); + out_buf.clear(); + prompt_embd_buf.clear(); + prompt_batch.reset(); + n_prompt = 0; + prompt_pos = 0; + step_idx = 0; + eos_step = -1; + chunks.clear(); + chunk_idx = 0; + n_voice_pos = 0; + chunk_budget = 0; + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + seq_id = inp->seq_id; + + if (!ensure_cache()) { + return 1; + } + + std::vector voice; + if (inp->speaker_ref) { + if (!encode_speaker(inp->speaker_ref, voice)) { + return 1; + } + } + + pack = pockettts_pack(info.model_variant); + + const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len), + pack.pad_short_text); + if (text.empty()) { + LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); + return 1; + } + + std::vector ids(text.size() + 16); + int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(), + (int32_t) ids.size(), false, false); + if (n_ids <= 0) { + LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n"); + return 1; + } + ids.resize((size_t) n_ids); + + // long inputs degrade badly, so each chunk restarts from the voice conditioning + // see split_into_best_sentences() in the reference + chunks = split_chunks(ids); + chunk_idx = 0; + if (chunks.size() > 1) { + LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size()); + } + + const int n_e = n_embd; + + // sequence order is voice, then text, then the audio BOS that starts generation + if (!voice.empty()) { + GGML_ASSERT(voice.size() % (size_t) n_e == 0); + if (bos_before_voice != LLAMA_TOKEN_NULL) { + push_embd_row(prompt_embd_buf, bos_before_voice); + } + prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); + } + // every later chunk rewinds to here and re-prompts, so the voice stays primed + n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e); + + for (llama_token t : chunks[0]) { + push_embd_row(prompt_embd_buf, t); + } + push_embd_row(prompt_embd_buf, audio_bos); + arm_chunk_budget(0); + + n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e); + prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e)); + prompt_batch->set_position_normal(0, seq_id); + prompt_pos = 0; + + seed = inp->seed; + out_type = inp->out_type; + + return 0; + } + + int32_t step_prompt(int32_t n_batch) override { + GGML_ASSERT(n_batch > 0); + if (prompt_pos >= n_prompt) { + return 0; + } + const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos); + llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch); + + if ((prompt_pos + n_tokens_batch) == n_prompt) { + batch_view.logits[n_tokens_batch - 1] = 1; + } + + if (llama_decode(lctx, batch_view) != 0) { + LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n"); + return -1; + } + + pos += n_tokens_batch; + prompt_pos += n_tokens_batch; + + if (prompt_pos >= n_prompt) { + prompt_batch.reset(); + prompt_embd_buf.clear(); + return 0; + } + return n_prompt - prompt_pos; + } + + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + (void) sampled; // the backbone output is continuous, there is no token to consume + + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.embd = const_cast(h_state_in); + // clip only reseeds when the seed changes, so pass the same one on every step + inp.seed = seed; + if (pack.temp > 0.0f) { + inp.temp = pack.temp; + } + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n"); + return 1; + } + if (out.is_eos && eos_step < 0) { + eos_step = step_idx; + } + // the frame of the stopping step is discarded, matching _autoregressive_generation(). + // the budget is the reference's fallback for a chunk whose eos head never fires + const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) || + step_idx >= chunk_budget; + if (chunk_done) { + if (eos_step < 0) { + LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx); + } + return finish_chunk(h_state_out, out_stop); + } + + feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats); + step_idx++; + if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) { + if (!flush_gen_wav()) { + return 1; + } + } + + decode_embd_batch batch_embd(const_cast(out.embd), 1, 1, n_embd); + batch_embd.set_position_normal(pos, seq_id); + batch_embd.batch.logits[0] = 1; + pos++; + + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: decode failed\n"); + return 1; + } + + const float * he = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(he, he + n_embd); + *h_state_out = h_state_buf.data(); + + return 0; + } + + int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override { + if (!flush_gen_wav()) { + return 1; + } + + *out_sample_rate = info.sample_rate; + if (out_n_samples) { + *out_n_samples = (int64_t) audio_pcm.size(); + } + + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *out_data = (const char *) audio_pcm.data(); + *out_data_len = audio_pcm.size() * sizeof(float); + return 0; + } + + out_buf.clear(); + if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) { + LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n"); + return 1; + } + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + +private: + bool ensure_cache() { + if (specials_ok) { + return true; + } + // bos_before_voice is optional, some packs do not insert it + bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>"); + audio_bos = find_special_token(vocab, "<|audio_bos|>"); + if (audio_bos == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n"); + return false; + } + const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr); + if (n_tok_embd == 0) { + LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n"); + return false; + } + tok_embd.resize(n_tok_embd); + if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) { + LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n"); + return false; + } + GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0); + specials_ok = true; + return true; + } + + // the table can be shorter than the vocab, so bound the row lookup + void push_embd_row(std::vector & dst, llama_token t) const { + const size_t n_rows = tok_embd.size() / (size_t) n_embd; + GGML_ASSERT(t >= 0 && (size_t) t < n_rows); + dst.insert(dst.end(), + tok_embd.begin() + (size_t) t * n_embd, + tok_embd.begin() + (size_t) (t + 1) * n_embd); + } + + // token ids of the pieces the reference splits on, see split_into_best_sentences(). + // the leading token is dropped, it is the tokenizer's dummy prefix + std::vector punct_ids(const char * s) const { + std::vector ids(16); + const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false); + if (n <= 1) { + return {}; + } + return std::vector(ids.begin() + 1, ids.begin() + n); + } + + // cut after runs of boundary tokens, so punctuation stays with the sentence it ends + static std::vector> split_on(const std::vector & ids, + const std::vector & boundary) { + std::vector> out; + size_t start = 0; + bool prev_was_boundary = false; + for (size_t i = 0; i < ids.size(); i++) { + const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end(); + if (!is_boundary && prev_was_boundary) { + out.emplace_back(ids.begin() + start, ids.begin() + i); + start = i; + } + prev_was_boundary = is_boundary; + } + out.emplace_back(ids.begin() + start, ids.end()); + return out; + } + + std::vector> split_chunks(const std::vector & ids) const { + if ((int) ids.size() <= max_chunk_tokens) { + return { ids }; + } + const std::vector eos_punct = punct_ids(".!...?"); + const std::vector mid_punct = punct_ids(",;:"); + + // oversized sentences are split again on weaker punctuation, else words get skipped + std::vector> segments; + for (auto & seg : split_on(ids, eos_punct)) { + if ((int) seg.size() <= max_chunk_tokens) { + segments.push_back(std::move(seg)); + continue; + } + auto sub = split_on(seg, mid_punct); + if (sub.size() > 1) { + for (auto & s : sub) { + segments.push_back(std::move(s)); + } + } else { + segments.push_back(std::move(seg)); + } + } + + std::vector> out; + for (auto & seg : segments) { + if (seg.empty()) { + continue; + } + if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) { + out.back().insert(out.back().end(), seg.begin(), seg.end()); + } else { + out.push_back(std::move(seg)); + } + } + if (out.empty()) { + out.push_back(ids); + } + for (const auto & c : out) { + if ((int) c.size() > max_chunk_tokens) { + LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, " + "generation may skip words\n", c.size(), max_chunk_tokens); + } + } + return out; + } + + // _estimate_max_gen_len() plus the per-chunk tail guess, both in frames + void arm_chunk_budget(size_t idx) { + const int n_tok = (int) chunks[idx].size(); + chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate); + // the pack may pin the tail, else the reference guesses it from the word count + frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3); + step_idx = 0; + eos_step = -1; + } + + // ends the current chunk and, if there is another, re-prompts it on top of the voice + int32_t finish_chunk(const float ** h_state_out, bool * out_stop) { + if (!flush_gen_wav()) { + return 1; + } + // the decoder restarts too, the next chunk's audio is not continuous with this one + dec_state.clear(); + + if (chunk_idx + 1 >= chunks.size()) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + chunk_idx++; + + // drop this chunk's text and audio, keep the voice conditioning + llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1); + pos = n_voice_pos; + + const int n_e = n_embd; + prompt_embd_buf.clear(); + for (llama_token t : chunks[chunk_idx]) { + push_embd_row(prompt_embd_buf, t); + } + push_embd_row(prompt_embd_buf, audio_bos); + arm_chunk_budget(chunk_idx); + + const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e); + GGML_ASSERT(n_rows > 0); + decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e); + batch.set_position_normal(pos, seq_id); + batch.batch.logits[n_rows - 1] = 1; + if (llama_decode(lctx, batch.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n"); + return 1; + } + pos += n_rows; + prompt_embd_buf.clear(); + + const float * he = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(he, he + n_embd); + *h_state_out = h_state_buf.data(); + *out_stop = false; + return 0; + } + + // same normalization as prepare_text_prompt() in the reference, it affects quality + static std::string prepare_text(const std::string & in, bool pad_short) { + std::string s; + s.reserve(in.size() + 1); + for (char c : in) { + if (c == '\n' || c == '\r') { + s += ' '; + } else if (c == ';') { + s += ','; + } else { + s += c; + } + } + const size_t b = s.find_first_not_of(' '); + const size_t e = s.find_last_not_of(' '); + if (b == std::string::npos) { + return ""; + } + s = s.substr(b, e - b + 1); + if (s[0] >= 'a' && s[0] <= 'z') { + s[0] = (char) (s[0] - 'a' + 'A'); + } + const unsigned char last = (unsigned char) s.back(); + if (std::isalnum(last)) { + s += '.'; + } + if (pad_short && count_words(s) < 5) { + s = std::string(8, ' ') + s; + } + return s; + } + + static int count_words(const std::string & s) { + int n = 0; + bool in_word = false; + for (char c : s) { + if (c == ' ') { + in_word = false; + } else if (!in_word) { + in_word = true; + n++; + } + } + return n; + } + + // runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame + bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { + continue; + } + if (mtmd_encode_chunk(mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n"); + break; + } + const float * embd = mtmd_get_output_embd(mctx); + const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk); + out.assign(embd, embd + n); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; + } + + // decodes the buffered latents, the mimi decoder state carries over between calls + bool flush_gen_wav() { + if (feats_buf.empty()) { + return true; + } + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; + inp.feats = feats_buf.data(); + inp.n_feats = feats_buf.size(); + inp.seed = seed; + inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data(); + inp.state_size = dec_state.size(); + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n"); + return false; + } + audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples); + dec_state.assign(out.state_data, out.state_data + out.state_size); + feats_buf.clear(); + return true; + } + + pockettts_pack_settings pack; + bool specials_ok = false; + llama_token bos_before_voice = LLAMA_TOKEN_NULL; + llama_token audio_bos = LLAMA_TOKEN_NULL; + std::vector tok_embd; + + llama_seq_id seq_id = 0; + int pos = 0; + std::vector prompt_embd_buf; + std::unique_ptr prompt_batch; + int n_prompt = 0; + int prompt_pos = 0; + uint32_t seed = UINT32_MAX; + // end-of-speech is latched, then a few more frames are generated as tail padding + int step_idx = 0; + int eos_step = -1; + int frames_after_eos = 3; + static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference + static constexpr double frame_rate = 12.5; + std::vector> chunks; + size_t chunk_idx = 0; + int n_voice_pos = 0; // KV positions held by the voice conditioning + int chunk_budget = 0; + + // latents are decoded a window at a time, the decoder state bridges the windows + size_t window_frames = 8; + std::vector feats_buf; + std::vector dec_state; + std::vector audio_pcm; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; +}; + static std::unique_ptr make_pipeline(llama_context * lctx, mtmd_context * mctx) { switch (mtmd_gen_audio_get_info(mctx).type) { case MTMD_GEN_AUDIO_TYPE_QWEN3TTS: return std::unique_ptr(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_POCKETTTS: + return std::unique_ptr(new pockettts_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } @@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n } int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled, - const float * h_state_in, const float ** h_state_out) { + const float * h_state_in, const float ** h_state_out, + bool * out_stop) { if (!ctx->pipeline) { return 1; } - return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out); + bool stop = false; + const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop); + if (out_stop) { + *out_stop = stop; + } + return ret; } int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate, diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7e5cf9b5098c..832f7171ac71 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp { mtmd_bitmap * speaker_ref; // optional, can be NULL const char * lang; // optional, can be NULL - int32_t top_k; - float top_p; + int32_t top_k; + float top_p; + uint32_t seed; // UINT32_MAX for random (default: random) enum mtmd_helper_gen_audio_outtype out_type; }; @@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt( int32_t n_batch); // generates one frame; must only be called after step_prompt() has returned 0 -// h_state_out is valid until next step_gen() or reset() call +// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token +// out_stop (optional) is set on end-of-speech, the caller must then stop the loop +// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated MTMD_API int32_t mtmd_helper_gen_audio_step_gen( mtmd_helper_gen_audio * ctx, llama_token sampled, const float * h_state_in, - const float ** h_state_out); + const float ** h_state_out, + bool * out_stop); // out_data valid until next get_output() or reset() call // out_n_samples (optional, can be NULL) receives the number of generated PCM samples @@ -261,8 +265,8 @@ struct gen_audio { int32_t step_prompt(int32_t n_batch) { return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch); } - int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) { - return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out); + int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) { + return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop); } int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) { return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 82b73d5cd126..4b9c45d62677 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -477,6 +477,7 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio std::vector gen_out_codes; // this frame's 16 sampled codes (GEN_CODE) + std::vector gen_out_feats; // this frame's continuous features, if any (GEN_CODE) std::vector gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV) std::vector gen_out_state; // state to feed into the next GEN_WAV call @@ -979,6 +980,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + audio_preproc = std::make_unique(ctx_a); + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } @@ -1798,16 +1803,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { // mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { - mtmd_gen_audio_info info; + mtmd_gen_audio_info info{}; + info.model_variant = ""; if (!ctx->ctx_gen_a) { info.type = MTMD_GEN_AUDIO_TYPE_NONE; return info; } + info.model_variant = clip_get_hparams(ctx->ctx_gen_a)->gen_model_variant.c_str(); switch (clip_get_projector_type(ctx->ctx_gen_a)) { case PROJECTOR_TYPE_QWEN3TTS_GEN: info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; info.sample_rate = 24000; break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { return info; } +mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { + mtmd_gen_inp inp{}; + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.seed = UINT32_MAX; + if (!ctx->ctx_gen_a) { + return inp; + } + + switch (clip_get_projector_type(ctx->ctx_gen_a)) { + case PROJECTOR_TYPE_QWEN3TTS_GEN: + // https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json + inp.top_k = 50; + inp.top_p = 1.0f; + inp.temp = 0.9f; // TODO: handle this on graph + break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + // https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py + inp.top_k = 50; + inp.top_p = 1.0f; + inp.temp = 0.7f; + break; + default: + break; + } + return inp; +} + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; if (!ctx_clip) { @@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 1; } + *out = {}; + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); @@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector out_embd(n_embd); std::vector out_codes; + std::vector out_feats; + bool is_eos = false; clip_encode_params params; - params.imgs = &batch; - params.n_threads = ctx->n_threads; - params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; - params.out_embd = &out_embd; - params.out_codes = &out_codes; - params.code0 = inp->code0; - params.top_k = inp->top_k; - params.top_p = inp->top_p; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; + params.out_embd = &out_embd; + params.out_codes = &out_codes; + params.out_feats = &out_feats; + params.code0 = inp->code0; + params.top_k = inp->top_k; + params.top_p = inp->top_p; + params.seed = inp->seed; + params.temp = inp->temp; + params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__); @@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in ctx->gen_out_embd = std::move(out_embd); ctx->gen_out_codes = std::move(out_codes); - - out->embd = ctx->gen_out_embd.data(); - out->codes = ctx->gen_out_codes.data(); - out->n_codes = ctx->gen_out_codes.size(); + ctx->gen_out_feats = std::move(out_feats); + + out->embd = ctx->gen_out_embd.data(); + out->codes = ctx->gen_out_codes.data(); + out->n_codes = ctx->gen_out_codes.size(); + out->feats = ctx->gen_out_feats.data(); + out->n_feats = ctx->gen_out_feats.size(); + out->is_eos = is_eos; return 0; } // MTMD_GEN_PROCESS_TYPE_GEN_WAV - if (!inp->codes || inp->n_codes == 0) { - LOG_ERR("%s: codes required for gen_wav\n", __func__); + const bool has_codes = inp->codes && inp->n_codes > 0; + const bool has_feats = inp->feats && inp->n_feats > 0; + if (has_codes == has_feats) { + LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__); return 1; } - std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector in_codes; + std::vector in_feats; + if (has_codes) { + in_codes.assign(inp->codes, inp->codes + inp->n_codes); + } else { + in_feats.assign(inp->feats, inp->feats + inp->n_feats); + } std::vector in_state; if (inp->state_data) { in_state.assign(inp->state_data, inp->state_data + inp->state_size); @@ -1885,7 +1943,10 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.imgs = &batch; params.n_threads = ctx->n_threads; params.gen_process = CLIP_GEN_PROCESS_GEN_WAV; - params.codes = &in_codes; + // gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation + params.seed = inp->seed; + params.codes = has_codes ? &in_codes : nullptr; + params.feats = has_feats ? &in_feats : nullptr; params.out_audio = &ctx->gen_out_audio; params.state_in = inp->state_data ? &in_state : nullptr; params.state_out = &ctx->gen_out_state; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index e5063d114cc9..c1a5921db2f3 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -344,18 +344,25 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, + MTMD_GEN_AUDIO_TYPE_POCKETTTS, }; + struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts + const char * model_variant; // name of the weight variant, can be nullptr if not applicable }; + MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); + enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.) MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio // for qwen3tts, this is code2wav + // for pocket-tts, this is mimi decoder }; + struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -364,21 +371,30 @@ struct mtmd_gen_inp { float * embd; // the hidden state from backbone, must have n_text_embd elements int32_t top_k; float top_p; + uint32_t seed; // UINT32_MAX for random + float temp; // sampling temperature, or noise scale for flow-matching decoders // for MTMD_GEN_PROCESS_TYPE_GEN_WAV + // pass either codes (discrete) or feats (continuous), depending on the pipeline int32_t * codes; size_t n_codes; + const float * feats; + size_t n_feats; const char * state_data; size_t state_size; }; + struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call // for MTMD_GEN_PROCESS_TYPE_GEN_CODE const int32_t * codes; - size_t n_codes; + size_t n_codes; + const float * feats; // continuous counterpart of codes + size_t n_feats; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + bool is_eos; // only set by pipelines having the EOS head inside mmproj // for MTMD_GEN_PROCESS_TYPE_GEN_WAV const float * audio; @@ -386,6 +402,10 @@ struct mtmd_gen_out { const char * state_data; size_t state_size; }; + +// defaults tuned for the loaded pipeline, callers override only what they care about +MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx); + // note: this API is stateless, caller must handle state management and audio frame accumulation MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, diff --git a/tools/tts/README.md b/tools/tts/README.md index dd84336c3990..1b08d5ef3218 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \ --tts-speaker-file speaker.mp3 \ --output out.wav ``` + +## Pocket TTS + +Available params: +- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it +- Note: `lang` is not used, the language is a property of the weights + +Example usage: + +```sh +llama-tts -m pocket-tts.gguf \ + -mm mmproj-pocket-tts.gguf \ + -p "Hello world" \ + --tts-speaker-file speaker.mp3 \ + --output out.wav +``` + +**Note for GGUF conversion:** + +The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/` directories, **not** the root directory: + +```sh +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf +``` diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 49c405e14784..fd7522f8defd 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -119,6 +119,7 @@ int main(int argc, char ** argv) { inp.lang = params.tts_lang.c_str(); inp.top_k = params.sampling.top_k; inp.top_p = params.sampling.top_p; + inp.seed = params.sampling.seed; inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; // @@ -143,8 +144,7 @@ int main(int argc, char ** argv) { } } - const llama_vocab * vocab = llama_model_get_vocab(model); - + // note: some pipelines ignore this token and use the hidden state instead auto sample_semantic_code = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); common_sampler_accept(smpl, t, true); @@ -159,19 +159,24 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { + bool stop = false; + while (!stop && n_frames < max_new) { const float * h_next = nullptr; // stage 2+3: semantic --> acoustic details --> audio waveform // step_gen() runs both stages and returns new h_state for next step - if (gen.step_gen(sampled, h_state, &h_next) != 0) { + if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) { LOG_ERR("step_gen failed at frame %d\n", n_frames); return 1; } + if (!h_next) { + break; // stopped without generating a frame + } + n_frames++; h_state = h_next; sampled = sample_semantic_code(); - timings.report(n_frames + 1); + timings.report(n_frames); } const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6; From cc078b45b635b3a59aa9adc1b888150ab67798a8 Mon Sep 17 00:00:00 2001 From: lnigam Date: Tue, 11 Aug 2026 18:46:26 +0530 Subject: [PATCH 015/211] Dflash support for nemotron-3.5 (#26905) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * conversion: skip untrained DFlash embeddings * Add Nemotron DFlash support * Add DFlash NVFP4 support * Address review comments * add missing output_s for nvfp4 * Include change for keeping residual for last layer also if requested in future dflash models * Update conversion/qwen.py Defensive check, not needed Co-authored-by: Sigbjørn Skjæret * Fixing bug introduced by merge conflict --------- Co-authored-by: Sigbjørn Skjæret --- conversion/base.py | 2 +- conversion/qwen.py | 11 ++++++++++- gguf-py/gguf/constants.py | 1 + src/llama-model.h | 5 +++-- src/models/dflash.cpp | 26 ++++++++++++++++---------- src/models/nemotron-h.cpp | 12 +++++++++++- 6 files changed, 42 insertions(+), 15 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 3572b77c21e2..718d5394495e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -829,7 +829,7 @@ def prepare_tensors(self): elif any(str(v.get("quant_algo")).endswith("NVFP4") for v in quant_layers.values() if isinstance(v, dict)): quant_algo = "NVFP4" - self._is_nvfp4 = quant_algo == "NVFP4" + self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4") self._is_mxfp4 = quant_method == "mxfp4" # NVFP4 weights are repacked and written directly to gguf_writer. diff --git a/conversion/qwen.py b/conversion/qwen.py index b4ae528bf2d4..ead435455dd6 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -647,10 +647,13 @@ def set_vocab(self): # own tokenizer logic, not the Qwen default). from . import get_model_class with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: - target_arch = json.load(f)["architectures"][0] + target_hparams = json.load(f) + target_arch = target_hparams["architectures"][0] target_cls = get_model_class(target_arch) if target_cls is not type(self): + if target_arch == "NemotronHForCausalLM": + setattr(self, "is_moe", "num_experts_per_tok" in target_hparams) target_cls.set_vocab(self) # ty: ignore[unresolved-attribute] else: super().set_vocab() @@ -688,6 +691,12 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca name = "model." + name return super().filter_tensors((name, gen)) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): + return + + yield from super().modify_tensors(data_torch, name, bid) + @ModelBase.register("Qwen3DSparkModel") class DSparkModel(DFlashModel): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a197cacd3fa5..98c4fa169143 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4726,6 +4726,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.D2T, ], MODEL_ARCH.DFLASH: [ + MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_Q, diff --git a/src/llama-model.h b/src/llama-model.h index 1dd0904387af..341cb66fbafe 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -623,8 +623,9 @@ struct llama_model { struct ggml_tensor * per_layer_model_proj = nullptr; struct ggml_tensor * per_layer_proj_norm = nullptr; - // eagle3 - struct ggml_tensor * fc = nullptr; // feature fusion layer + // eagle3 / dflash feature fusion layer + struct ggml_tensor * fc = nullptr; + struct ggml_tensor * fc_s = nullptr; struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping // dspark diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index daff6e78f1cf..bfbdb28ee5e9 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -79,6 +79,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -97,6 +98,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { } fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); + fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -205,7 +207,7 @@ template <> llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { ggml_tensor * cur = build_inp_embd_enc(); - cur = build_lora_mm(model.fc, cur); + cur = build_lora_mm(model.fc, cur, model.fc_s); cb(cur, "fc_out", -1); cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); @@ -460,9 +462,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cb(cur, "ffn_norm", il); cur = build_ffn(cur, - layer.ffn_up, NULL, NULL, - layer.ffn_gate, NULL, NULL, - layer.ffn_down, NULL, NULL, + layer.ffn_up, NULL, layer.ffn_up_s, + layer.ffn_gate, NULL, layer.ffn_gate_s, + layer.ffn_down, NULL, layer.ffn_down_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); @@ -479,15 +481,17 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_embd = cur; // lm_head from the target model (shared via ctx_other) - auto * output = model.output; + auto * output = model.output; + auto * output_s = model.output_s; if (output == nullptr) { GGML_ASSERT(cparams.ctx_other != nullptr); const auto * model_other = llama_get_model(cparams.ctx_other); GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); - output = model_other->output; + output = model_other->output; + output_s = model_other->output_s; } - cur = build_lora_mm(output, cur); + cur = build_lora_mm(output, cur, output_s); cb(cur, "result_output", -1); res->t_logits = cur; @@ -655,15 +659,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_ cb(cur, "result_norm", -1); // lm_head from the target model (shared via ctx_other) - auto * output = model.output; + auto * output = model.output; + auto * output_s = model.output_s; if (output == nullptr) { GGML_ASSERT(cparams.ctx_other != nullptr); const auto * model_other = llama_get_model(cparams.ctx_other); GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection"); - output = model_other->output; + output = model_other->output; + output_s = model_other->output_s; } - cur = build_lora_mm(output, cur); + cur = build_lora_mm(output, cur, output_s); cb(cur, "result_output", -1); res->t_logits = cur; diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index cd2af3179c3a..f02674c64610 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -177,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ auto * inp = build_inp_mem_hybrid(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]; for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + struct ggml_tensor * inpSA = inpL; // norm @@ -195,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_ffn_layer(cur, model, il); } - if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -209,6 +212,13 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ } cur = inpL; + if (extract_final_inp) { + res->t_layer_inp[n_layer] = cur; + + if (inp_out_ids && cparams.embeddings_nextn_masked) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + } cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); From 5d16e81dd9896355d36b363bbf786abaa8f6995f Mon Sep 17 00:00:00 2001 From: ynankani Date: Tue, 11 Aug 2026 13:19:05 +0000 Subject: [PATCH 016/211] convert : keep quantization scales for nemotron --mtp export (#26903) Signed-off-by: ynankani --- conversion/nemotron.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/conversion/nemotron.py b/conversion/nemotron.py index e5075020c1e3..c46cec143866 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -275,10 +275,18 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None elif cls.mtp_only: # --mtp: export the MTP head plus the tensors it shares with the target model + # Include lm_head scale sidecars so NVFP4 packing sees them. keep = name in ( "backbone.embeddings.weight", "backbone.norm_f.weight", "lm_head.weight", + "lm_head.weight_scale", + "lm_head.weight_scale_2", + "lm_head.weight_scale_inv", + "lm_head.input_scale", + "lm_head.input_global_scale", + "lm_head.weight_global_scale", + "lm_head.weight_packed", ) if not keep: return None From 2468576f241235452013308597e6de1b78866996 Mon Sep 17 00:00:00 2001 From: Niklas Wenzel Date: Tue, 11 Aug 2026 15:58:53 +0200 Subject: [PATCH 017/211] requirements: use stable torch packages on s390x (#26864) --- examples/model-conversion/requirements.txt | 2 +- requirements/requirements-convert_hf_to_gguf.txt | 6 +----- requirements/requirements-convert_lora_to_gguf.txt | 2 -- tools/mtmd/requirements.txt | 8 +------- 4 files changed, 3 insertions(+), 15 deletions(-) diff --git a/examples/model-conversion/requirements.txt b/examples/model-conversion/requirements.txt index 229b2ec75b75..d2cd357ec95c 100644 --- a/examples/model-conversion/requirements.txt +++ b/examples/model-conversion/requirements.txt @@ -1,6 +1,6 @@ --extra-index-url https://download.pytorch.org/whl/cpu torch -torchvision +torchvision; platform_machine != "s390x" transformers huggingface-hub accelerate diff --git a/requirements/requirements-convert_hf_to_gguf.txt b/requirements/requirements-convert_hf_to_gguf.txt index f80fdc1f6402..b1f7c863e27e 100644 --- a/requirements/requirements-convert_hf_to_gguf.txt +++ b/requirements/requirements-convert_hf_to_gguf.txt @@ -2,8 +2,4 @@ --extra-index-url https://download.pytorch.org/whl/cpu ## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility -torch==2.11.0; platform_machine != "s390x" - -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly -torch>=0.0.0.dev0; platform_machine == "s390x" +torch==2.11.0 diff --git a/requirements/requirements-convert_lora_to_gguf.txt b/requirements/requirements-convert_lora_to_gguf.txt index d091d564846b..5758076c41dc 100644 --- a/requirements/requirements-convert_lora_to_gguf.txt +++ b/requirements/requirements-convert_lora_to_gguf.txt @@ -1,4 +1,2 @@ -r ./requirements-convert_hf_to_gguf.txt --extra-index-url https://download.pytorch.org/whl/cpu -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly diff --git a/tools/mtmd/requirements.txt b/tools/mtmd/requirements.txt index f26d8e912a37..d646ca7b02f2 100644 --- a/tools/mtmd/requirements.txt +++ b/tools/mtmd/requirements.txt @@ -2,11 +2,5 @@ --extra-index-url https://download.pytorch.org/whl/cpu pillow~=11.3.0 -## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility -torch==2.11.0; platform_machine != "s390x" # check_requirements: ignore "==" +torch==2.11.0 # check_requirements: ignore "==" torchvision==0.26.0; platform_machine != "s390x" # check_requirements: ignore "==" - -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly -torch>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "==" -torchvision>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "==" From 38406d597fe651aafc7d843d270cb25b5ec2617a Mon Sep 17 00:00:00 2001 From: Bartowski <3266127+bartowski1182@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:18:19 -0400 Subject: [PATCH 018/211] imatrix.cpp: Move finite check and only check touched experts (#26861) --- tools/imatrix/imatrix.cpp | 73 ++++++++++++++++++++++++--------------- 1 file changed, 46 insertions(+), 27 deletions(-) diff --git a/tools/imatrix/imatrix.cpp b/tools/imatrix/imatrix.cpp index 3431a4eca84b..f5fee621840a 100644 --- a/tools/imatrix/imatrix.cpp +++ b/tools/imatrix/imatrix.cpp @@ -222,6 +222,15 @@ static void compute_cossim(std::vector & tstats) { } } +static bool all_finite(const float * v, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (!std::isfinite(v[i])) { + return false; + } + } + return true; +} + bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * user_data) { GGML_UNUSED(user_data); @@ -299,33 +308,39 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * exit(1); //GGML_ABORT("fatal error"); } LOG_DBGV(2, "%s[%d]: %32s, %s, %5d x %5d, %d\n", __func__, m_last_chunk, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[2], (int)src1->type); - // loop over all possible experts, regardless if they are used or not in the batch - for (int64_t ex = 0; ex < n_as; ++ex) { - size_t e_start = ex*src1->ne[0]; - - for (int64_t idx = 0; idx < n_ids; ++idx) { - for (int64_t row = 0; row < src1->ne[2]; ++row) { - const int excur = *(const int32_t *) (m_ids.data() + row*ids->nb[1] + idx*ids->nb[0]); - GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check + const int64_t ne0 = src1->ne[0]; + const int64_t n_tokens = src1->ne[2]; - if (excur != ex) continue; + // single pass over the routing ids + std::vector touched(n_as, 0); + for (int64_t idx = 0; idx < n_ids; ++idx) { + for (int64_t row = 0; row < n_tokens; ++row) { + const int32_t ex = *(const int32_t *) (m_ids.data() + row * ids->nb[1] + idx * ids->nb[0]); - const int64_t i11 = idx % src1->ne[1]; - const int64_t i12 = row; - const float * x = (const float *)(data + i11*src1->nb[1] + i12*src1->nb[2]); + GGML_ASSERT(ex >= 0 && ex < n_as); // sanity check - e.counts[ex]++; + const int64_t i11 = idx % src1->ne[1]; + const float * x = (const float *) (data + i11 * src1->nb[1] + row * src1->nb[2]); + float * acc = e.values.data() + ex * ne0; - for (int64_t j = 0; j < src1->ne[0]; ++j) { - e.values[e_start + j] += x[j] * x[j]; - if (!std::isfinite((float)e.values[e_start + j])) { - LOG_ERR("%f detected in %s\n", (float)e.values[e_start + j], wname.c_str()); - exit(1); - } - } + e.counts[ex]++; + touched[ex] = 1; + for (int64_t j = 0; j < ne0; ++j) { + acc[j] += x[j] * x[j]; } } + } + + // check for non-finite values, only checking experts that were routed to and touched + for (int64_t ex = 0; ex < n_as; ++ex) { + if (touched[ex] && !all_finite(e.values.data() + ex * ne0, ne0)) { + LOG_ERR("%s: non-finite values detected in %s\n", __func__, wname.c_str()); + exit(1); + } + } + + for (int64_t ex = 0; ex < n_as; ++ex) { const int32_t n_chunk = e.counts[ex] / chunk_size; if (n_chunk > m_last_chunk) { const int32_t chunk_step = n_chunk - m_last_chunk; @@ -366,24 +381,28 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * } LOG_DBGV(2, "%s[%d]: %32s, %s, %5d x %5d x %5d, %d\n", __func__, m_last_chunk, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[1], (int)src1->ne[2], (int)src1->type); + const int64_t ne0 = src1->ne[0]; + for (int64_t i3 = 0; i3 < src1->ne[3]; ++i3) { for (int64_t i2 = 0; i2 < src1->ne[2]; ++i2) { // handle 3D+ tensors, but flatten 3D+ activations when model tensor is 2D const int64_t mat_id = (i3 % src0->ne[3]) * src0->ne[2] + (i2 % src0->ne[2]); - const int64_t mat_start = mat_id * src1->ne[0]; + float * acc = e.values.data() + mat_id * ne0; for (int64_t row = 0; row < src1->ne[1]; ++row) { const float * x = (const float *) (data + row * src1->nb[1] + i2 * src1->nb[2] + i3 * src1->nb[3]); - for (int64_t j = 0; j < src1->ne[0]; ++j) { - e.values[mat_start + j] += x[j] * x[j]; - if (!std::isfinite((float)e.values[j])) { - LOG_ERR("%f detected in %s\n", (float)e.values[j], wname.c_str()); - exit(1); - } + for (int64_t j = 0; j < ne0; ++j) { + acc[j] += x[j] * x[j]; } } } } + + // check for non-finite values + if (!all_finite(e.values.data(), e.values.size())) { + LOG_ERR("%s: non-finite values detected in %s\n", __func__, wname.c_str()); + exit(1); + } // only 1 count in practice, except when a tensor is used for both MUL_MAT_ID and MUL_MAT for (size_t i = 0; i < e.counts.size(); ++i) { e.counts[i] += ggml_nrows(src1) / n_mat; From 70dfba5aee36793fb51ae649723b3c30ed9e99d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Tue, 11 Aug 2026 17:48:27 +0200 Subject: [PATCH 019/211] ci : add windows-rocm to check-release (#26897) [no release] --- .github/workflows/release.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 3a48a57c1b68..7f47c72135f8 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -749,6 +749,7 @@ jobs: name: llama-bin-win-cpu-${{ matrix.arch }}.zip windows-rocm: + needs: [check-release] runs-on: windows-2022 strategy: From ba360efe1f574ebae727aad64112d18ecedca85a Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Tue, 11 Aug 2026 10:58:54 -0500 Subject: [PATCH 020/211] chat : tighten bare function parsing for Qwen models (#26793) --- common/chat.cpp | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/common/chat.cpp b/common/chat.cpp index 6cbf23b50508..01053ddde804 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1166,6 +1166,16 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ data.prompt += data.generation_prompt; } + std::vector tool_call_starts = { "" }; + + // Match complete opener for Qwen3-Coder models that occasionally omit the + // starting . The model may hallucinate a tool name, but it is preferable over + // constraining on + foreach_function(inputs.tools, [&](const json & tool) { + const std::string name = tool.at("function").at("name"); + tool_call_starts.push_back(""); + }); + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { auto generation_prompt = p.literal(GEN_PREFIX); @@ -1238,7 +1248,7 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1)); return generation_prompt + - (reasoning << p.content(p.until_one_of({ "", "" }, - // Trigger on " Date: Tue, 11 Aug 2026 19:52:12 +0300 Subject: [PATCH 021/211] spec : update speculative-simple (#26904) * spec : update speculative-simple * cont : simplify * cont : clean-up --- common/speculative.cpp | 80 ++--------- common/speculative.h | 6 - examples/speculative-simple/README.md | 47 ++++++- .../speculative-simple/speculative-simple.cpp | 130 +++++++++--------- tools/server/server-context.cpp | 7 +- 5 files changed, 118 insertions(+), 152 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 2ee1e6b84818..0082e5fc5d3d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -171,12 +171,6 @@ struct common_speculative_impl { // (optional) serialize/restore per-seq internal state (e.g. eagle3's deferred boundary). virtual bool get_state(llama_seq_id /*seq_id*/, std::vector & /*data*/) const { return false; } virtual void set_state(llama_seq_id /*seq_id*/, const std::vector & /*data*/) {} - - // true if this implementation requires the target context to extract post-norm embeddings - virtual bool need_embd() const = 0; - - // true if this implementation requires the target context to extract pre-norm embeddings - virtual bool need_embd_nextn() const { return false; } }; struct common_speculative_impl_draft_simple : public common_speculative_impl { @@ -193,6 +187,10 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; auto * ctx_tgt = this->params.ctx_tgt; + if (!ctx_dft) { + throw std::runtime_error("draft-simple requires a draft context"); + } + SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n"); SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", @@ -385,10 +383,6 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; @@ -907,10 +901,6 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { pending_g_last[seq_id].resize(n_embd_dec); std::memcpy(pending_g_last[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd_dec * sizeof(float)); } - - bool need_embd() const override { - return false; - } }; // DFlash: block-diffusion drafting with a draft-side KV cache injection @@ -1247,10 +1237,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_draft_mtp : public common_speculative_impl { @@ -1689,14 +1675,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const size_t row_bytes = (size_t) n_embd * sizeof(float); std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); } - - bool need_embd() const override { - return false; - } - - bool need_embd_nextn() const override { - return true; - } }; // state of self-speculation (simple implementation, not ngram-map) @@ -1743,10 +1721,6 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_map_k : public common_speculative_impl { @@ -1801,10 +1775,6 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { common_ngram_map_accept(config[seq_id], n_accepted); } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_mod : public common_speculative_impl { @@ -1980,10 +1950,6 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } } } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_cache : public common_speculative_impl { @@ -2123,10 +2089,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative { @@ -2322,7 +2284,6 @@ common_speculative_init_result::common_speculative_init_result( const bool spec_mtp = std::find(params.speculative.types.begin(), params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - GGML_ASSERT(has_draft || spec_mtp); auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -2560,34 +2521,6 @@ bool common_speculative_process(common_speculative * spec, const llama_batch & b return result; } -bool common_speculative_need_embd(common_speculative * spec) { - if (spec == nullptr) { - return false; - } - - for (auto & impl : spec->impls) { - if (impl->need_embd()) { - return true; - } - } - - return false; -} - -bool common_speculative_need_embd_nextn(common_speculative * spec) { - if (spec == nullptr) { - return false; - } - - for (auto & impl : spec->impls) { - if (impl->need_embd_nextn()) { - return true; - } - } - - return false; -} - void common_speculative_draft(common_speculative * spec) { if (spec == nullptr) { return; @@ -2672,7 +2605,10 @@ void common_speculative_draft(common_speculative * spec) { void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, uint16_t n_accepted) { common_speculative_impl * impl = spec->impl_last[seq_id]; - GGML_ASSERT(impl); + if (impl == nullptr) { + GGML_ASSERT(n_accepted == 0); + return; + } { common_time_meas tm(impl->t_accept_us, !impl->gen_perf); diff --git a/common/speculative.h b/common/speculative.h index c6986affdadf..06b0992ed681 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -67,12 +67,6 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co // process the batch and update the internal state of the speculative context bool common_speculative_process(common_speculative * spec, const llama_batch & batch); -// true if any implementation requires target post-norm embeddings to be extracted -bool common_speculative_need_embd(common_speculative * spec); - -// true if any implementation requires target nextn embeddings to be extracted -bool common_speculative_need_embd_nextn(common_speculative * spec); - // generate drafts for the sequences specified with `common_speculative_get_draft_params` void common_speculative_draft(common_speculative * spec); diff --git a/examples/speculative-simple/README.md b/examples/speculative-simple/README.md index f72129b3f92e..b81583f00bcc 100644 --- a/examples/speculative-simple/README.md +++ b/examples/speculative-simple/README.md @@ -3,10 +3,47 @@ Demonstration of basic greedy speculative decoding ```bash +# spec-type draft-simple ./bin/llama-speculative-simple \ - -m ../models/qwen2.5-32b-coder-instruct/ggml-model-q8_0.gguf \ - -md ../models/qwen2.5-1.5b-coder-instruct/ggml-model-q4_0.gguf \ - -f test.txt -c 0 -ngl 99 --color on \ - --sampling-seq k --top-k 1 -fa on --temp 0.0 \ - -ngld 99 --spec-draft-n-max 16 --spec-draft-n-draft-min 5 --draft-p-min 0.9 + -hf ggml-org/Qwen3-8B-Base-GGUF:Q8_0 \ + -hfd ggml-org/Qwen3-0.6B-Base-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-simple --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-mtp +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3.6-27B-GGUF:Q8_0 \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-mtp --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-mtp (with shared KV cache) +# note: this model needs a token at the start to somewhat work without the chat template +./bin/llama-speculative-simple \ + -hf ggml-org/Gemma-4-31B-it-GGUF:Q8_0 \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-mtp --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-eagle3 +./bin/llama-speculative-simple \ + -hf ggml-org/gpt-oss-20b-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-eagle3 --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-dflash +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3-8B-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-dflash --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-dspark +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3-8B-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-dspark --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 ``` diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index c727e8139d9c..487ae03abfa7 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -51,48 +51,23 @@ int main(int argc, char ** argv) { const llama_vocab * vocab = llama_model_get_vocab(model_tgt); - // load the draft model - llama_model_ptr model_dft; - llama_context_ptr ctx_dft; + // load the draft model (if any) - this also creates the MTP draft context when MTP speculation is enabled + common_speculative_init_result_ptr spec_init; - // TODO: simplify this logic { - const auto & params_spec = params.speculative.draft; + common_params params_dft = common_base_params_to_speculative(params); - auto params_dft = params; - - params_dft.n_outputs_max = params.n_parallel; - params_dft.n_outputs_max_per_seq = 1; - - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - - if (params_spec.cpuparams.n_threads > 0) { - params_dft.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; - params_dft.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads; - } - - params_dft.tensor_buft_overrides = params.speculative.draft.tensor_buft_overrides; - - auto mparams_dft = common_model_params_to_llama(params_dft); - - model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft)); - if (model_dft == nullptr) { - LOG_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); - return 1; - } - - auto cparams = common_context_params_to_llama(params_dft); - ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); + spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); params.speculative.draft.ctx_tgt = ctx_tgt; - params.speculative.draft.ctx_dft = ctx_dft.get(); + params.speculative.draft.ctx_dft = spec_init->context(); } + llama_context * ctx_dft = params.speculative.draft.ctx_dft; + // check if the context supports partial sequence removal - const bool use_ckpt_tgt = (common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); - const bool use_ckpt_dft = (common_context_can_seq_rm(ctx_dft.get()) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); + const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; if (use_ckpt_tgt) { LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n"); @@ -138,9 +113,30 @@ int main(int argc, char ** argv) { // target model sampling context common_sampler_ptr smpl(common_sampler_init(model_tgt, params.sampling)); - // eval the prompt - llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); - llama_decode(ctx_dft.get(), llama_batch_get_one(inp.data(), inp.size() - 1)); + // init the speculator + const auto & params_spec = params.speculative; + + struct common_speculative * spec = common_speculative_init(params.speculative, 1); + + if (spec == nullptr) { + LOG_ERR("%s", "failed to initialize speculative decoding\n"); + return 1; + } + + // eval the prompt on the target and feed it to the speculative implementation(s) + { + llama_batch batch_prompt = llama_batch_init(inp.size(), 0, 1); + for (size_t i = 0; i < inp.size() - 1; ++i) { + common_batch_add(batch_prompt, inp[i], i, { seq_id }, false); + } + + llama_decode(ctx_tgt, batch_prompt); + + if (!common_speculative_process(spec, batch_prompt)) { + LOG_ERR("%s", "failed to process speculative prompt\n"); + return 1; + } + } // note: keep the last token separate! llama_token id_last = inp.back(); @@ -151,18 +147,12 @@ int main(int argc, char ** argv) { int n_past = inp.size() - 1; - // init the speculator - const auto & params_spec = params.speculative; - - struct common_speculative * spec = common_speculative_init(params.speculative, 1); - common_speculative_begin(spec, seq_id, prompt_tgt); llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); - size_t n_draft = 0; - llama_tokens draft; + common_prompt_checkpoint ckpt; const auto t_enc_end = ggml_time_us(); @@ -184,13 +174,20 @@ int main(int argc, char ** argv) { llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id)); if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } + // determine the max draft that fits the remaining context and generation budget + int n_draft_max = (int) llama_n_ctx(ctx_tgt) - n_past - 2; + if (params.n_predict >= 0) { + n_draft_max = std::min(n_draft_max, params.n_predict - n_predict - 1); + } + n_draft_max = std::max(n_draft_max, 0); + // generate a new draft common_speculative_get_draft_params(spec, seq_id) = { /* .drafting = */ true, - /* .n_max = */ -1, + /* .n_max = */ n_draft_max, /* .n_past = */ n_past, /* .id_last = */ id_last, /* .prompt = */ &prompt_tgt, @@ -198,9 +195,6 @@ int main(int argc, char ** argv) { }; common_speculative_draft(spec); - // save the original draft size - n_draft = draft.size(); - // save a checkpoint of the target context before evaluating the draft // this allows us to restore the state if partial draft acceptance occurs if (!draft.empty()) { @@ -209,10 +203,13 @@ int main(int argc, char ** argv) { } } - { - ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + // reset the draft context to the checkpoint before verification + if (ctx_dft) { + if (use_ckpt_dft) { + ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1); + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1); } } else { // we have a previous (partial) draft to reuse from checkpoint restoration @@ -236,10 +233,10 @@ int main(int argc, char ** argv) { llama_decode(ctx_tgt, batch_tgt); } - // evaluate the same batch with the draft model - { - // TODO: extend to support MTP, Eagle, etc. See server code for reference - llama_decode(ctx_dft.get(), batch_tgt); + // feed the batch to the speculative implementation(s) - this drives the draft model, MTP, Eagle3, etc. + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("%s", "failed to process speculative batch\n"); + break; } // only save the sampler sampler state if we use checkpoints @@ -248,6 +245,9 @@ int main(int argc, char ** argv) { smpl_save.reset(common_sampler_clone(smpl.get())); } + // save the size of the draft being verified + const size_t n_draft = draft.size(); + // sample from the full target batch and return the accepted tokens based on the target sampler // // for each token to be accepted, the sampler would have to sample that same token @@ -264,8 +264,8 @@ int main(int argc, char ** argv) { // check for partial draft acceptance: // if the context doesn't support partial sequence removal, restore the checkpoint // and make the accepted tokens the new partial draft for the next iteration - if (use_ckpt_tgt && ids.size() - 1 < draft.size()) { - LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, draft.size()); + if (use_ckpt_tgt && ids.size() - 1 < n_draft) { + LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft); draft = std::move(ids); @@ -275,10 +275,10 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1); } - { - ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (ctx_dft) { + ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1); + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1); } prompt_tgt.resize(ckpt.n_tokens); @@ -329,8 +329,11 @@ int main(int argc, char ** argv) { { LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past); - llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1); - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, n_past, -1); + llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1); + + if (ctx_dft) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, n_past, -1); + } } if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { @@ -356,6 +359,7 @@ int main(int argc, char ** argv) { LOG_INF("\n"); LOG_INF("draft:\n\n"); + common_speculative_print_stats(spec); LOG_INF("\n"); LOG_INF("target:\n\n"); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 3b5f6a12185c..d75c6856dfa9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -397,12 +397,7 @@ struct server_slot { bool need_embd() const { GGML_ASSERT(task); - return task->need_embd() || (spec && common_speculative_need_embd(spec)); - } - - bool need_embd_nextn() const { - GGML_ASSERT(task); - return spec && common_speculative_need_embd_nextn(spec); + return task->need_embd(); } // if the context does not have a memory module then all embeddings have to be computed within a single ubatch From 5988633170f66d71570275b89f5c78e0c236e67a Mon Sep 17 00:00:00 2001 From: 0 <1939455790@qq.com> Date: Wed, 12 Aug 2026 01:46:23 +0800 Subject: [PATCH 022/211] cuda : add warp-per-row wkv7 kernel for single-token decode (#26111) --- ggml/src/ggml-cuda/wkv.cu | 56 +++++++++++++++++++++++++++++++++++++- tests/test-backend-ops.cpp | 1 + 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu index d2fced705e09..2361112124f0 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu @@ -141,6 +141,57 @@ static __global__ void rwkv_wkv7_f32(const int B, const int T, const int C, cons } } +template +static __global__ void __launch_bounds__(WARP_SIZE * rows_per_block, 2) +rwkv_wkv7_f32_t1_warp_row(const int T, const int C, const int H, const float * r, const float * w, const float * k, const float * v, const float * a, const float * b, const float * s, float * dst) { + constexpr int head_size = CUDA_WKV_BLOCK_SIZE; + constexpr int half_head = head_size / 2; + + const int lane = threadIdx.x; + const int row = blockIdx.y * rows_per_block + threadIdx.y; + const int bid = blockIdx.x; + + const int batch_i = bid / H; + const int head_i = bid % H; + const int state_size = C * head_size; + const int head_off = head_i * head_size; + const int t = batch_i * C + head_off + row; + + __shared__ float _r[head_size], _w[head_size], _k[head_size], _a[head_size], _b[head_size]; + + if (threadIdx.y == 0) { + _r[lane] = r[batch_i * C + head_off + lane]; + _w[lane] = w[batch_i * C + head_off + lane]; + _k[lane] = k[batch_i * C + head_off + lane]; + _a[lane] = a[batch_i * C + head_off + lane]; + _b[lane] = b[batch_i * C + head_off + lane]; + + _r[lane + half_head] = r[batch_i * C + head_off + lane + half_head]; + _w[lane + half_head] = w[batch_i * C + head_off + lane + half_head]; + _k[lane + half_head] = k[batch_i * C + head_off + lane + half_head]; + _a[lane + half_head] = a[batch_i * C + head_off + lane + half_head]; + _b[lane + half_head] = b[batch_i * C + head_off + lane + half_head]; + } + __syncthreads(); + + const int64_t state_base = batch_i * state_size + head_i * head_size * head_size + row * head_size; + const float s0 = s[state_base + lane]; + const float s1 = s[state_base + lane + half_head]; + const float sa = warp_reduce_sum(_a[lane] * s0 + _a[lane + half_head] * s1); + + const float vt = v[t]; + const float st0 = s0 * _w[lane] + _k[lane] * vt + sa * _b[lane]; + const float st1 = s1 * _w[lane + half_head] + _k[lane + half_head] * vt + sa * _b[lane + half_head]; + const float y = warp_reduce_sum(st0 * _r[lane] + st1 * _r[lane + half_head]); + + dst[T * C + state_base + lane] = st0; + dst[T * C + state_base + lane + half_head] = st1; + + if (lane == 0) { + dst[t] = y; + } +} + void ggml_cuda_op_rwkv_wkv6(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float * k_d = (const float *)dst->src[0]->data; const float * v_d = (const float *)dst->src[1]->data; @@ -191,7 +242,10 @@ void ggml_cuda_op_rwkv_wkv7(ggml_backend_cuda_context & ctx, ggml_tensor * dst) GGML_ASSERT(C % H == 0); GGML_ASSERT(C / H == CUDA_WKV_BLOCK_SIZE || C / H == CUDA_WKV_BLOCK_SIZE * 2); - if (C / H == CUDA_WKV_BLOCK_SIZE) { + if (T / B == 1 && C / H == CUDA_WKV_BLOCK_SIZE) { + constexpr int rows_per_block = 4; + rwkv_wkv7_f32_t1_warp_row<<>>(T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d); + } else if (C / H == CUDA_WKV_BLOCK_SIZE) { rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d); } else { rwkv_wkv7_f32<<>>(B, T, C, H, r_d, w_d, k_d, v_d, a_d, b_d, s_d, dst_d); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a4c22156c63c..86dfabbf5b3b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8827,6 +8827,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 128, 4)); test_cases.emplace_back(new test_rwkv_wkv7(GGML_TYPE_F32, 32, 64, 1, 1)); + test_cases.emplace_back(new test_rwkv_wkv7(GGML_TYPE_F32, 32, 64, 1, 4)); test_cases.emplace_back(new test_rwkv_wkv7(GGML_TYPE_F32, 32, 64, 32, 1)); test_cases.emplace_back(new test_rwkv_wkv7(GGML_TYPE_F32, 32, 64, 32, 4)); test_cases.emplace_back(new test_rwkv_wkv7(GGML_TYPE_F32, 32, 64, 128, 4)); From ebb546b7e961bd46fd9ed0387ffd14ca86b6fe1b Mon Sep 17 00:00:00 2001 From: Rafail Giavrimis <47496212+grafail@users.noreply.github.com> Date: Tue, 11 Aug 2026 18:50:03 +0100 Subject: [PATCH 023/211] CUDA: only disable CUDA graphs when mul_mat_id actually needs a stream sync (#26802) --- ggml/src/ggml-cuda/ggml-cuda.cu | 39 ++++++++++++++++++++++++++++----- tests/test-backend-ops.cpp | 4 ++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1d4f4dfbd214..cb7e9330c8c6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1865,6 +1865,37 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); } +// returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization +// [TAG_MUL_MAT_ID_CUDA_GRAPHS] +static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return true; + } + + if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + if (ggml_is_quantized(src0->type)) { + if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + return false; + } + } else if (GGML_CUDA_CC_IS_AMD(cc)) { + return false; + } + } + + if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) { + return false; + } + + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + return false; + } + + return true; +} + static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -1907,7 +1938,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization - // TODO: add asserts to verify this. should work with CUDA, HIP, etc. + GGML_ASSERT(ggml_cuda_mul_mat_id_needs_sync(dst, cc)); cudaStream_t stream = ctx.stream(); GGML_ASSERT(nb12 % nb11 == 0); @@ -2522,10 +2553,8 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (node->op == GGML_OP_MUL_MAT_ID) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); - if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { - // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs - // TODO: figure out a way to enable for larger batch sizes, without hurting performance + if (ggml_cuda_mul_mat_id_needs_sync(node, cc)) { + // the mul_mat_id fallback path synchronizes the stream, so we cannot use CUDA graphs // ref: https://github.com/ggml-org/llama.cpp/pull/18958 use_cuda_graph = false; #ifndef NDEBUG diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 86dfabbf5b3b..c27f0b92e814 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9012,6 +9012,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 128, k, {12,1}, {1,1})); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 16, false, 50, 200, k)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 16, true, 50, 200, k)); } @@ -9043,6 +9045,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 2, 2, b, 32, 8192, 64)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); } test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1)); From 7b13a8404d7e219c13d1a243e2a21a857a6e99d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Tue, 11 Aug 2026 20:20:40 +0200 Subject: [PATCH 024/211] ci : add missing release check (#26923) --- .github/workflows/release.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7f47c72135f8..af34c9499396 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -750,6 +750,8 @@ jobs: windows-rocm: needs: [check-release] + if: ${{ needs.check-release.outputs.should_release == 'true' }} + runs-on: windows-2022 strategy: From 0b1bad14ff204627636aeb1de22ddcd5acb859d4 Mon Sep 17 00:00:00 2001 From: ruanslv Date: Tue, 11 Aug 2026 16:15:20 -0400 Subject: [PATCH 025/211] chat : fix muse-glimmer detection of tool calls after EOM (#26879) * chat : fix muse-glimmer swallowing a trailing tool call into content Muse Glimmer routinely answers the user and calls a tool in a single generation. The template terminates a message with <|eom|> when more messages follow in the same turn and <|eot|> only at the end of the turn, so the answer is closed by <|eom|> and the call opens a fresh header: <|eom|><|start|>assistant to=<|message|>... The final-message rule read content with until("<|eot|>"), which assumed the user-facing message is always last. There is no <|eot|> before the call, so content ran to the end of the turn, absorbed the markup, and no tool_calls were emitted - the tool never ran. On a tau2-bench telecom run this hit 43 turns across 19 of 114 tasks. Stop the answer at <|eom|> and parse what follows as tool calls. Adds models/templates/muse-glimmer.jinja and four parser tests: a plain answer, the <|eom|> junction, markup quoted in an answer staying content, and tool markup inside the to=self channel staying reasoning. * address comment --- common/chat.cpp | 6 +- models/templates/muse-glimmer.jinja | 211 ++++++++++++++++++++++++++++ tests/test-chat.cpp | 46 ++++++ 3 files changed, 261 insertions(+), 2 deletions(-) create mode 100644 models/templates/muse-glimmer.jinja diff --git a/common/chat.cpp b/common/chat.cpp index 01053ddde804..faf51dcd27d7 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -3155,7 +3155,8 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat auto analysis = p.ref("analysis"); auto recipient = p.optional(p.literal(" to=user")); - auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + p.content(p.until("<|eot|>"))); + auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + + p.content(p.until_one_of({ "<|eot|>", "<|eom|>" }))); if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { auto string_value = p.ac( @@ -3211,7 +3212,8 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { return p.zero_or_more(start + analysis) + start + tool_calls; } - return p.zero_or_more(start + analysis) + start + (tool_calls | final_msg); + auto trailing_calls = p.optional(p.literal("<|eom|>") + start + tool_calls); + return p.zero_or_more(start + analysis) + start + (tool_calls | (final_msg + trailing_calls)); } return p.zero_or_more(start + analysis) + start + final_msg; diff --git a/models/templates/muse-glimmer.jinja b/models/templates/muse-glimmer.jinja new file mode 100644 index 000000000000..7507f3c9f388 --- /dev/null +++ b/models/templates/muse-glimmer.jinja @@ -0,0 +1,211 @@ +{# + Template: Muse Glimmer ATEM Chat Template + Renders the ATEM tool-calling protocol: reasoning channel (to=self), tool + channels (to=), and the user channel, plus tool definitions and the + valid-recipient list in the system block. + + Whitespace note: every tag uses the {%- -%} / {{- -}} stripping markers, so + the indentation below is purely for readability and contributes nothing to + the rendered output. +#} +{%- macro render_content(content) -%} + {%- if content is string -%} + {{- content -}} + {%- elif content is not none -%} + {%- for part in content -%} + {%- if part['type'] == 'image' -%} + {{- '<|patch|>' -}} + {%- elif part['type'] == 'video' -%} + {{- '<|video|>' -}} + {%- elif part['type'] == 'text' -%} + {{- part['text'] -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} +{%- endmacro -%} +{%- macro render_atem(tc) -%} + {%- set args = tc.function.arguments -%} + {%- if args is not mapping -%} + {{- raise_exception('Muse Glimmer ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.') -}} + {%- endif -%} + {{- '\n\n' -}} + {%- for k, v in args.items() -%} + {{- '' -}} + {%- if v is boolean -%} + {%- if v -%} + true + {%- else -%} + false + {%- endif -%} + {%- elif v is none -%} + null + {%- elif v is mapping or (v is iterable and v is not string) -%} + {{- v | tojson -}} + {%- else -%} + {{- v -}} + {%- endif -%} + {{- '\n' -}} + {%- endfor -%} + {{- '\n' -}} +{%- endmacro -%} +{%- macro render_tool_defs(tools) -%} + {{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}} + {{- 'You can invoke a function by writing a "" block like the following:\n' -}} + {{- '\n\n$PARAMETER_VALUE\n...\n\n\n\n' -}} + {{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}} + {{- 'Here are the functions available in JSONSchema format:\n' -}} + {{- '// Tool metadata\n' -}} + {%- set nsns = namespace(seen=[]) -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split('.')[0] -%} + {%- if tns not in nsns.seen -%} + {%- set nsns.seen = nsns.seen + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%} + {%- for tns in nsns.seen -%} + {{- '{"name": ' + (tns | tojson) + ', "description": ' + ((nd[tns] if tns in nd else '') | tojson) + '}\n' -}} + {%- endfor -%} + {{- '// Function schemas' -}} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {{- '\n{"name": ' + (fn.name | tojson) + ', "description": ' + (fn.description | tojson) + ', "parameters": ' + (fn.parameters | tojson) + '}' -}} + {%- endfor -%} + {{- '\n\nHere\'s an example of how to call a function in the tool set:\n' -}} + {{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than `example_tool_name.example_function_name`)\n\n' -}} + {{- 'to=example_tool_name.example_function_name\n\n' -}} + {{- '\n\n' -}} + {{- 'value_1\n' -}} + {{- 'This is the value for the second parameter\nthat can span\n"multiple" lines\n\n' -}} + {{- '\n' -}} +{%- endmacro -%} +{%- macro render_reasoning() -%} + {%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else 'high' -%} + {{- 'Reasoning strength: ' + rs + '.' -}} +{%- endmacro -%} +{%- macro render_system_meta(tools) -%} + {%- set rns = namespace(recipients=['"self"'], nslist=[]) -%} + {%- if tools -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split('.')[0] -%} + {%- if tns not in rns.nslist -%} + {%- set rns.nslist = rns.nslist + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- for tns in rns.nslist -%} + {%- set rns.recipients = rns.recipients + ['"' + tns + '.*"'] -%} + {%- endfor -%} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ['"user"'] -%} + {{- '# Valid recipients: ' + rns.recipients | join(', ') + '.' -}} +{%- endmacro -%} +{{- bos_token -}} +{%- set ns = namespace(has_system=false) -%} +{%- for m in messages -%} + {%- if m['role'] == 'system' -%} + {%- set ns.has_system = true -%} + {%- endif -%} +{%- endfor -%} +{%- if not ns.has_system -%} + {{- '<|start|>system<|message|>You are a helpful AI assistant.' -}} + {%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else '2026-01-04' -%} + {{- '\nKnowledge cutoff: ' + kc + '.' -}} + {%- if current_date is defined and current_date -%} + {{- '\nCurrent date: ' + current_date + '.' -}} + {%- elif strftime_now is defined -%} + {{- '\nCurrent date: ' + strftime_now('%Y-%m-%d') + '.' -}} + {%- endif -%} + {{- '\n\n' -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- '\n\n' -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- '\n\n' -}} + {{- render_system_meta(tools) -}} + {{- '<|eot|>' -}} +{%- endif -%} +{%- for message in messages -%} + {%- set role = message['role'] -%} + {%- set end_token = '<|eom|>' if (not loop.last and messages[loop.index0 + 1]['role'] == role) else '<|eot|>' -%} + {%- if role == 'system' -%} + {#- Callers sometimes write the directive into the system prompt themselves. + Normalise "Reasoning effort" to "Reasoning strength" (jinja has no + case-insensitive replace, hence the four realistic casings), then skip + the kwarg-driven line below if the prompt already carries one. -#} + {%- set sys_text = render_content(message['content']) + | replace('Reasoning effort', 'Reasoning strength') + | replace('Reasoning Effort', 'Reasoning Strength') + | replace('reasoning effort', 'reasoning strength') + | replace('REASONING EFFORT', 'REASONING STRENGTH') -%} + {{- '<|start|>system<|message|>' -}} + {{- sys_text -}} + {%- if 'reasoning strength' not in (sys_text | lower) -%} + {{- '\n\n' -}} + {{- render_reasoning() -}} + {%- endif -%} + {%- if tools -%} + {{- '\n\n' -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- '\n\n' -}} + {{- render_system_meta(tools) -}} + {{- '<|eot|>' -}} + {%- elif role == 'user' -%} + {{- '<|start|>user<|message|>' -}} + {{- render_content(message['content']) -}} + {{- '<|eot|>' -}} + {%- elif role == 'tool' -%} + {%- set tname = message.get('name') -%} + {%- if not tname -%} + {%- set tcid = message.get('tool_call_id') -%} + {%- set rns = namespace(name=tcid if tcid else '') -%} + {%- for m in messages -%} + {%- if m.get('tool_calls') -%} + {%- for tc in m['tool_calls'] -%} + {%- if tcid is not none and tc.id is defined and tc.id == tcid -%} + {%- set rns.name = tc.function.name -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- set tname = rns.name -%} + {%- endif -%} + {{- '<|start|>tool ' + tname + '<|message|>\n' -}} + {{- render_content(message['content']) -}} + {{- '\n<|eot|>' -}} + {%- elif role == 'assistant' -%} + {%- if message.get('reasoning_content') -%} + {{- '<|start|>assistant to=self<|message|>' + message['reasoning_content'] + '<|eom|>' -}} + {%- endif -%} + {%- if message.get('tool_calls') -%} + {%- for tc in message['tool_calls'] -%} + {{- '<|start|>assistant to=' + tc.function.name + '<|message|>' -}} + {{- render_atem(tc) -}} + {%- if loop.last -%} + {{- end_token -}} + {%- else -%} + {{- '<|eom|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {%- set recipient = message.get('recipient') or 'user' -%} + {%- set end_turn = message.get('end_turn') -%} + {%- if end_turn is none -%} + {%- set end_turn = not (recipient and recipient != 'user') -%} + {%- endif -%} + {{- '<|start|>assistant' -}} + {%- if recipient -%} + {{- ' to=' + recipient -}} + {%- endif -%} + {{- '<|message|>' -}} + {{- render_content(message['content']) -}} + {{- ('<|eot|>' if end_turn else '<|eom|>') -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|start|>assistant' -}} +{%- endif -%} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index f54a58f9b67c..3cf81ca8e73d 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -5843,6 +5843,52 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // Muse Glimmer format tests + { + auto tst = peg_tester("models/templates/muse-glimmer.jinja", detailed_debug); + + const std::string call_markup = + "\n" + "\n" + "1\n" + "\n" + ""; + + // A plain answer is unaffected + tst.test(" to=user<|message|>Hello, world!\nWhat's up?<|eot|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist) + .run(); + + // "Inform then act": the model answers the user and calls a tool in ONE generation, + // closing the answer with <|eom|>. The answer must stop there rather than swallow it. + tst.test(" to=user<|message|>Hello, world!\nWhat's up?<|eom|>" + "<|start|>assistant to=special_function<|message|>" + + call_markup) + .tools({ special_function_tool }) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_with_content_and_tool_call("Hello, world!\nWhat's up?", "special_function", + "{\"arg1\":1}")) + .run(); + + // Markup quoted in an answer has no preceding <|eom|>, so it stays content instead of + // becoming an invocation the user never asked for + tst.test(" to=user<|message|>You invoke it like this:\n" + call_markup + "<|eot|>") + .tools({ special_function_tool }) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect_content("You invoke it like this:\n" + call_markup) + .run(); + + // Tool markup inside the analysis channel is reasoning, not a call + tst.test(" to=self<|message|>I could use " + call_markup + " here<|eom|>" + "<|start|>assistant to=user<|message|>Hello!<|eot|>") + .tools({ special_function_tool }) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect_reasoning("I could use " + call_markup + " here") + .expect_content("Hello!") + .run(); + } + // GPT-OSS format tests { auto tst = peg_tester("models/templates/openai-gpt-oss-120b.jinja", detailed_debug); From cb27fe9c35d953b49b49a0048e73decdac0338c6 Mon Sep 17 00:00:00 2001 From: lhez Date: Tue, 11 Aug 2026 22:02:28 -0700 Subject: [PATCH 026/211] opencl: use flat mv q5_k when weight exceeds image1d_buffer_t limit (#26880) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 19dca4ced0b8..5d0324dad0ca 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7076,6 +7076,19 @@ inline bool enable_adreno_trans_weight(const ggml_backend_opencl_context *backen return ((elem_num < 128 * 1024 * 1024) && adreno_kernel && shape_ok); // max element num: 2**27 } +inline bool enable_adreno_trans_weight_q5_K(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) { + if (!use_adreno_kernels(backend_ctx, tensor)) { + return false; + } + + const size_t elem_num = ggml_nelements(tensor); + const size_t q_img_width = elem_num / 8; + const size_t qh_img_width = elem_num / 16; + + return q_img_width <= backend_ctx->image_max_buffer_size && + qh_img_width <= backend_ctx->image_max_buffer_size; +} + static inline bool use_flat_gemv_for_large_m_q4_K(const ggml_tensor *tensor) { // gemv_noshuffle variant perf drops for large M, use flat variant for large M. // threshold is well above typical hidden/FFN dims, but below typical vocab sizes. @@ -9255,7 +9268,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, #ifdef GGML_OPENCL_USE_ADRENO_KERNELS cl_kernel kernel = backend_ctx->kernel_convert_block_q5_K; - if (use_adreno_kernels(backend_ctx, tensor)) { + if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) { kernel = backend_ctx->kernel_convert_block_q5_K_noshuffle; } #else @@ -9290,7 +9303,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, tensor->extra = extra; #ifdef GGML_OPENCL_USE_ADRENO_KERNELS - if (use_adreno_kernels(backend_ctx, tensor)) { + if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) { int M = tensor->ne[1]; int K = tensor->ne[0]; @@ -10388,7 +10401,7 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer, CL_CHECK(clReleaseMemObject(data_device)); return; } - if (use_adreno_kernels(backend_ctx, tensor)) { + if (enable_adreno_trans_weight_q5_K(backend_ctx, tensor)) { int M = tensor->ne[1]; int K = tensor->ne[0]; @@ -18928,7 +18941,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co } // q5_K x fp32 - if (src0t == GGML_TYPE_Q5_K && src1t == GGML_TYPE_F32) { + if (src0t == GGML_TYPE_Q5_K && src1t == GGML_TYPE_F32 && + enable_adreno_trans_weight_q5_K(backend_ctx, src0)) { ggml_cl_mul_mat_q5_K_f32_adreno(backend, src0, src1, dst); return; } From 6eff59326205854cf4b9edd45e67ee3436c16d58 Mon Sep 17 00:00:00 2001 From: Wang Zhiyu Date: Wed, 12 Aug 2026 13:05:13 +0800 Subject: [PATCH 027/211] convert : handle per_layer_config in Gemma4 (transformers 5.15) (#26882) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: handle nested global_head_dim in Gemma4 config Gemma-4 E4B models have global_head_dim inside text_config rather than at the top level. Add fallback to support both layouts. * fix: add fallback for global_head_dim to support per_layer_config format * fix: read head_dim only from full_attention layers in per_layer_config and num_global_key_value_heads compatibility * fix: added fallback for num_global_key_value_heads * fix: read per_layer_config from root hparams * fix: delete unused text_config * cleanup and fixes --------- Co-authored-by: Sigbjørn Skjæret --- conversion/gemma.py | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/conversion/gemma.py b/conversion/gemma.py index c552df732b0f..f15a10a38bb6 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -665,7 +665,18 @@ def set_gguf_parameters(self): swa_layers = [t == "sliding_attention" for t in self.hparams["layer_types"]] self.gguf_writer.add_sliding_window_pattern(swa_layers) - head_dim_full = self.hparams["global_head_dim"] + per_layer_config = self.hparams.get("per_layer_config") + layer_types = self.hparams.get("layer_types", []) + if (head_dim_full := self.hparams.get("global_head_dim")) is None and per_layer_config is not None: + for layer_idx, layer_config in per_layer_config.items(): + layer_idx = int(layer_idx) + if layer_idx < len(layer_types): + if layer_types[layer_idx] == "full_attention" and "head_dim" in layer_config: + head_dim_full = layer_config["head_dim"] + break + + assert head_dim_full is not None + head_dim_swa = self.hparams["head_dim"] # correct the head dim for global/swa layers self.gguf_writer.add_key_length(head_dim_full) @@ -685,8 +696,14 @@ def set_gguf_parameters(self): n_ff_arr = [n_ff if il < first_kv_shared_layer_idx else n_ff * 2 for il in range(self.block_count)] self.gguf_writer.add_feed_forward_length(n_ff_arr) - # handle num_global_key_value_heads - num_key_value_heads_full = self.hparams.get("num_global_key_value_heads") + if (num_key_value_heads_full := self.hparams.get("num_global_key_value_heads")) is None and per_layer_config is not None: + for layer_idx, layer_config in per_layer_config.items(): + layer_idx = int(layer_idx) + if layer_idx < len(layer_types): + if layer_types[layer_idx] == "full_attention" and "num_key_value_heads" in layer_config: + num_key_value_heads_full = layer_config["num_key_value_heads"] + break + num_key_value_heads_swa = self.hparams.get("num_key_value_heads") if num_key_value_heads_full is not None and num_key_value_heads_swa is not None: value_arr = [num_key_value_heads_swa if is_swa else num_key_value_heads_full for is_swa in swa_layers] @@ -708,7 +725,19 @@ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: # IMPORTANT: this ROPE_FREQS tensor is ONLY used by the full_attention layers rope_params_full = self.hparams["rope_parameters"]["full_attention"] assert rope_params_full["rope_type"] == "proportional" - head_dim_full = (self.hparams["global_head_dim"]) + + per_layer_config = self.hparams.get("per_layer_config") + if (head_dim_full := self.hparams.get("global_head_dim")) is None and per_layer_config is not None: + layer_types = self.hparams.get("layer_types", []) + for layer_idx, layer_config in per_layer_config.items(): + layer_idx = int(layer_idx) + if layer_idx < len(layer_types): + if layer_types[layer_idx] == "full_attention" and "head_dim" in layer_config: + head_dim_full = layer_config["head_dim"] + break + + assert head_dim_full is not None + partial_rotary_factor_full = rope_params_full["partial_rotary_factor"] n_rot_full = int(head_dim_full * partial_rotary_factor_full / 2) n_unrot_full = int(head_dim_full / 2) - n_rot_full From 55f453b92452c9a103ae963543addce02e800f90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?O=C4=9Fuzhan=20Akkaya?= Date: Wed, 12 Aug 2026 01:06:16 -0400 Subject: [PATCH 028/211] wavtokenizer-dec : bound posnet/convnext block_count against n_layer_all (#26892) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * wavtokenizer-dec : bound posnet/convnext block_count against n_layer_all * Update src/llama-model.cpp Co-authored-by: Sigbjørn Skjæret --------- Co-authored-by: Sigbjørn Skjæret --- src/llama-model.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0e27cb41713b..c810055050da 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1122,6 +1122,9 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, hparams.convnext.n_embd); ml.get_key(LLM_KV_CONVNEXT_BLOCK_COUNT, hparams.convnext.n_layer); + + GGML_ASSERT(hparams.posnet.n_layer <= hparams.n_layer_all); + GGML_ASSERT(hparams.convnext.n_layer <= hparams.n_layer_all); } GGML_ASSERT(hparams.n_expert <= LLAMA_MAX_EXPERTS); From a7cd2f0e98ffa4e939276287be1fa76891908d2e Mon Sep 17 00:00:00 2001 From: michaeltrabalka-tech Date: Wed, 12 Aug 2026 01:07:23 -0400 Subject: [PATCH 029/211] vulkan: add TQ2_0 (ternary) support (#25850) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * vulkan: TQ2_0 (ternary) support — dequant + dedicated mul_mat_vec + matmul via dequant_funcs First Vulkan ternary type in ggml. Correctness: OM-125m TQ2_0 vs F16 top-12 logprobs identical to 4 decimals fully offloaded (float dequant path, no Q8_K activation quant). Speed at 125m ~= F16 (overhead-bound at this scale); the bandwidth win targets larger BitNet SKUs. MMQ/int-dot path intentionally not wired yet. Co-Authored-By: Claude Fable 5 * tests: enable TQ2_0 in backend-ops type lists Vulkan now implements TQ2_0 (dequant, mul_mat_vec, mul_mm, get_rows); backends without support skip via not-supported as usual. TQ1_0 stays disabled. Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Michael Trabalka Co-authored-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 23 ++++ .../vulkan-shaders/dequant_funcs.glsl | 14 +++ .../vulkan-shaders/dequant_funcs_cm2.glsl | 41 +++++++ .../vulkan-shaders/dequant_tq2_0.comp | 31 ++++++ .../vulkan-shaders/mul_mat_vec_tq2_0.comp | 102 ++++++++++++++++++ .../vulkan-shaders/mul_mm_funcs.glsl | 16 +++ .../src/ggml-vulkan/vulkan-shaders/types.glsl | 24 +++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 3 +- tests/test-backend-ops.cpp | 6 +- 9 files changed, 257 insertions(+), 3 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq2_0.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq2_0.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 45fa97f81297..c815d4ff99b5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4627,6 +4627,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ2_0], matmul_tq2_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) @@ -4667,6 +4668,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) @@ -4739,6 +4741,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); @@ -4783,6 +4786,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); @@ -4873,6 +4877,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -4921,6 +4926,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); @@ -4968,6 +4974,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); @@ -5047,6 +5054,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0].f32acc, matmul_tq2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -5094,6 +5102,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_subgroup_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_subgroup_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_subgroup_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_subgroup_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_subgroup_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_subgroup_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_subgroup_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); @@ -5123,6 +5132,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); @@ -5226,6 +5236,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32", arr_dmmv_q5_1_f32_f32_len[reduc], arr_dmmv_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32", arr_dmmv_q8_0_f32_f32_len[reduc], arr_dmmv_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32", arr_dmmv_q2_k_f32_f32_len[reduc16], arr_dmmv_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f32_f32", arr_dmmv_tq2_0_f32_f32_len[reduc16], arr_dmmv_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32", arr_dmmv_q3_k_f32_f32_len[reduc16], arr_dmmv_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32", arr_dmmv_q5_k_f32_f32_len[reduc16], arr_dmmv_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5253,6 +5264,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32", arr_dmmv_q5_1_f16_f32_len[reduc], arr_dmmv_q5_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32", arr_dmmv_q8_0_f16_f32_len[reduc], arr_dmmv_q8_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32", arr_dmmv_q2_k_f16_f32_len[reduc16], arr_dmmv_q2_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f16_f32", arr_dmmv_tq2_0_f16_f32_len[reduc16], arr_dmmv_tq2_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32", arr_dmmv_q3_k_f16_f32_len[reduc16], arr_dmmv_q3_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32", arr_dmmv_q4_k_f16_f32_len[reduc16], arr_dmmv_q4_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32", arr_dmmv_q5_k_f16_f32_len[reduc16], arr_dmmv_q5_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5307,6 +5319,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", arr_dmmv_id_q5_1_f32_f32_len[reduc], arr_dmmv_id_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", arr_dmmv_id_q8_0_f32_f32_len[reduc], arr_dmmv_id_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", arr_dmmv_id_q2_k_f32_f32_len[reduc16], arr_dmmv_id_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ2_0], "mul_mat_vec_id_tq2_0_f32", arr_dmmv_id_tq2_0_f32_f32_len[reduc16], arr_dmmv_id_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", arr_dmmv_id_q3_k_f32_f32_len[reduc16], arr_dmmv_id_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", arr_dmmv_id_q4_k_f32_f32_len[reduc16], arr_dmmv_id_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", arr_dmmv_id_q5_k_f32_f32_len[reduc16], arr_dmmv_id_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5368,6 +5381,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); @@ -5396,6 +5410,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_K], "get_rows_q2_k", get_rows_q2_k_len, get_rows_q2_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ2_0], "get_rows_tq2_0", get_rows_tq2_0_len, get_rows_tq2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q3_K], "get_rows_q3_k", get_rows_q3_k_len, get_rows_q3_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_K], "get_rows_q4_k", get_rows_q4_k_len, get_rows_q4_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_K], "get_rows_q5_k", get_rows_q5_k_len, get_rows_q5_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5424,6 +5439,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_K], "get_rows_q2_k_f32", get_rows_q2_k_f32_len, get_rows_q2_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ2_0], "get_rows_tq2_0_f32", get_rows_tq2_0_f32_len, get_rows_tq2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q3_K], "get_rows_q3_k_f32", get_rows_q3_k_f32_len, get_rows_q3_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_K], "get_rows_q4_k_f32", get_rows_q4_k_f32_len, get_rows_q4_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_K], "get_rows_q5_k_f32", get_rows_q5_k_f32_len, get_rows_q5_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -7638,6 +7654,7 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return nullptr; @@ -7712,6 +7729,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return nullptr; @@ -7781,6 +7799,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return nullptr; @@ -7874,6 +7893,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return nullptr; @@ -7946,6 +7966,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return nullptr; @@ -18014,6 +18035,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: break; default: return false; @@ -18119,6 +18141,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_TQ2_0: case GGML_TYPE_I32: return true; default: diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl index d902ff3a67bd..627932bd3547 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl @@ -608,6 +608,20 @@ vec2 get_dm(uint ib, uint a_offset) { } #endif +#if defined(DATA_A_TQ2_0) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + // elem e -> byte qs[(e/128)*32 + e%32], bits 2*((e%128)/32); w = q - 1 (d applied via get_dm) + const uint qsi = (iqs / 128) * 32 + (iqs % 32); // iqs even -> qsi, qsi+1 in same group/level + const uint shift = 2 * ((iqs % 128) / 32); + + const uvec2 qs = uvec2(data_a[a_offset + ib].qs[qsi], data_a[a_offset + ib].qs[qsi + 1]); + return vec2((qs >> shift) & 3) - 1.0; +} +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), 0); +} +#endif + #if defined(DATA_A_Q3_K) vec2 dequantize(uint ib, uint iqs, uint a_offset) { iqs /= 2; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl index 6bf2cb0e08ed..46cc69cb26ed 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl @@ -247,6 +247,44 @@ f16vec4 dequantFuncQ8_0_v(const in decodeBufQ8_0 bl, const in uint blockCoords[2 return f16vec4(vec4(qi) * vec4(float(d))); } +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0 { + block_tq2_0 block; +}; + +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0_packed16 { + block_tq2_0_packed16 block; +}; + +float16_t dequantFuncTQ2_0(const in decodeBufTQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + decodeBufTQ2_0_packed16 bl16 = decodeBufTQ2_0_packed16(bl); + const uint idx = coordInBlock[1]; + + const uint qsshift = (idx & 0x60) >> 4; // 0,2,4,6 + + uint qs = uint32_t(bl16.block.qs[((idx & 0x80) >> 3) + ((idx & 0x1E) >> 1)]); + qs = (qs >> qsshift) & 0x0303; + qs = unpack8(qs)[idx & 1]; + + return bl.block.d * (float16_t(int(qs)) - float16_t(1.0)); +} + +f16vec4 dequantFuncTQ2_0_v(const in decodeBufTQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + + const uint qsshift = (idx & 0x60) >> 4; // 0,2,4,6 + const uint qsi = ((idx & 0x80) >> 2) + (idx & 0x1C); // byte index of 4-aligned group + + const uint qsw = (uint(bl.block.qs[qsi])) + | (uint(bl.block.qs[qsi + 1]) << 8) + | (uint(bl.block.qs[qsi + 2]) << 16) + | (uint(bl.block.qs[qsi + 3]) << 24); + const u8vec4 q = unpack8((qsw >> qsshift) & 0x03030303); + + return bl.block.d * (f16vec4(q) - f16vec4(1.0)); +} + layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufQ2_K { block_q2_K block; }; @@ -1368,6 +1406,9 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords #elif defined(DATA_A_Q8_0) #define dequantFuncA dequantFuncQ8_0 #define dequantFuncA_v dequantFuncQ8_0_v +#elif defined(DATA_A_TQ2_0) +#define dequantFuncA dequantFuncTQ2_0 +#define dequantFuncA_v dequantFuncTQ2_0_v #elif defined(DATA_A_Q2_K) #define dequantFuncA dequantFuncQ2_K #define dequantFuncA_v dequantFuncQ2_K_v diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq2_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq2_0.comp new file mode 100644 index 000000000000..9475c9a23897 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq2_0.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) { + const uint i = gl_WorkGroupID.x * 256 + wgy; + if (i >= p.nel / QUANT_K) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint ip = tid / 32; // group 0,1 (128 elems each) + const uint il = tid - 32 * ip; // byte in group 0..31 + + const uint y_idx = i * QUANT_K + 128 * ip + il; + + const uint8_t qs = data_a[i].qs[32 * ip + il]; + + const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d); + data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(int((qs >> 0) & 3) - 1)); + data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(int((qs >> 2) & 3) - 1)); + data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(int((qs >> 4) & 3) - 1)); + data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(int((qs >> 6) & 3) - 1)); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq2_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq2_0.comp new file mode 100644 index 000000000000..689cfc42a51e --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq2_0.comp @@ -0,0 +1,102 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require + +#include "mul_mat_vec_base.glsl" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + +// ternary TQ2_0: w = (q - 1) * d. Same qs group/level layout as q2_K, but a +// single f16 scale per 256-block and no mins: +// sum_e b_e * (q_e - 1) * d = d * (sum_e b_e * q_e - sum_e b_e) +void calc_superblock(const uint a_offset, const uint b_offset, const uint v_im, const uint q_offset, const uint y_offset, const uint i, const uint num_blocks_per_row, const uint first_row, const uint num_rows) { + const uint y_idx = i * QUANT_K + y_offset; + + [[unroll]] for (uint n = 0; n < num_rows; ++n) { + const uint ib0 = a_offset + (first_row+n)*num_blocks_per_row; + if (i >= num_blocks_per_row) { + continue; + } + + const uint32_t qs_u32 = uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2]) | (uint32_t(data_a_packed16[ib0 + i].qs[q_offset / 2 + 8]) << 16); + const vec4 qs_u32_0 = vec4(unpack8(qs_u32 & 0x03030303)); + const vec4 qs_u32_2 = vec4(unpack8((qs_u32 >> 2) & 0x03030303)); + const vec4 qs_u32_4 = vec4(unpack8((qs_u32 >> 4) & 0x03030303)); + const vec4 qs_u32_6 = vec4(unpack8((qs_u32 >> 6) & 0x03030303)); + + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d); + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + vec2 b0 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 0]); + vec2 b16 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 8]); + vec2 b32 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 16]); + vec2 b48 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 24]); + vec2 b64 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 32]); + vec2 b80 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 40]); + vec2 b96 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 48]); + vec2 b112 = vec2(data_b_v2[(j*p.batch_stride_b + b_offset + y_idx) / 2 + 56]); + + FLOAT_TYPE sumq = FLOAT_TYPE(0.0); + FLOAT_TYPE sumb = FLOAT_TYPE(0.0); + [[unroll]] for (int l = 0; l < 2; ++l) { + sumq = fma(FLOAT_TYPE(b0[l]), FLOAT_TYPE(qs_u32_0[l ]), + fma(FLOAT_TYPE(b16[l]), FLOAT_TYPE(qs_u32_0[l+2]), + fma(FLOAT_TYPE(b32[l]), FLOAT_TYPE(qs_u32_2[l ]), + fma(FLOAT_TYPE(b48[l]), FLOAT_TYPE(qs_u32_2[l+2]), + fma(FLOAT_TYPE(b64[l]), FLOAT_TYPE(qs_u32_4[l ]), + fma(FLOAT_TYPE(b80[l]), FLOAT_TYPE(qs_u32_4[l+2]), + fma(FLOAT_TYPE(b96[l]), FLOAT_TYPE(qs_u32_6[l ]), + fma(FLOAT_TYPE(b112[l]), FLOAT_TYPE(qs_u32_6[l+2]), sumq)))))))); + sumb += FLOAT_TYPE(b0[l]) + FLOAT_TYPE(b16[l]) + FLOAT_TYPE(b32[l]) + FLOAT_TYPE(b48[l]) + + FLOAT_TYPE(b64[l]) + FLOAT_TYPE(b80[l]) + FLOAT_TYPE(b96[l]) + FLOAT_TYPE(b112[l]); + } + temp[j][n] = fma(d, sumq - sumb, temp[j][n]); + } + } +} + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads are used to process each block + const uint it_size = gl_WorkGroupSize.x/16; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid%16; // 0...15 + const uint ix = tid/16; + + const uint v_im = itid/8; // 0 or 1. 0 computes 0..., 1 computes 128... + const uint v_in = itid - 8*v_im; // 0...7 + + const uint l0 = 2*v_in; // 0...15 + const uint q_offset = 32*v_im + l0; + const uint y_offset = 128*v_im + l0; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + for (uint i0 = 0; i0 < num_blocks_per_row; i0 += it_size) + calc_superblock(a_offset, b_offset, v_im, q_offset, y_offset, i0 + ix, num_blocks_per_row, first_row, num_rows); + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 31dfefec8f94..63af2ce68574 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -182,6 +182,22 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); +#elif defined(DATA_A_TQ2_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = (idx % 128) * 2; // elem 0,2,4..254 + + const uint qsi = (iqs / 128) * 32 + (iqs % 32); // byte pair start + const uint shift = 2 * ((iqs % 128) / 32); // 0,2,4,6 + + const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); + const float d = float(data_a[ib].d); + + const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0); + + buf_a[buf_idx] = FLOAT_TYPEV2(v.xy); #elif defined(DATA_A_Q3_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl index 9616a26c7b39..adb1bb8b32b5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl @@ -303,6 +303,30 @@ struct block_q2_K_packed32 #define DATA_A_QUANT_K #endif +#define QUANT_K_TQ2_0 256 + +// ternary (BitNet): 2-bit codes, w = (q - 1) * d; qs layout matches q2_K's +// two 32-byte groups with four bit-levels per byte +struct block_tq2_0 +{ + uint8_t qs[QUANT_K_TQ2_0/4]; + float16_t d; +}; + +struct block_tq2_0_packed16 +{ + uint16_t qs[QUANT_K_TQ2_0/4/2]; + float16_t d; +}; + +#if defined(DATA_A_TQ2_0) +#define QUANT_K QUANT_K_TQ2_0 +#define QUANT_R 1 +#define A_TYPE block_tq2_0 +#define A_TYPE_PACKED16 block_tq2_0_packed16 +#define DATA_A_QUANT_K +#endif + #define QUANT_K_Q3_K 256 struct block_q3_K diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index c93d6eecee1b..6c9f76af1c9c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -72,6 +72,7 @@ const std::vector type_names = { "iq4_nl", "mxfp4", "nvfp4", + "tq2_0", "bf16", }; @@ -733,7 +734,7 @@ void process_shaders() { for (const auto& tname : type_names) { // mul mat vec std::string data_a_key = "DATA_A_" + to_uppercase(tname); - std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}})); string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c27f0b92e814..631a0cfcd553 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8000,7 +8000,8 @@ static const ggml_type all_types[] = { GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, - // GGML_TYPE_TQ1_0, GGML_TYPE_TQ2_0, // TODO: implement for all backends + GGML_TYPE_TQ2_0, + // GGML_TYPE_TQ1_0, // TODO: implement for all backends GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_XS, @@ -8027,7 +8028,8 @@ static const ggml_type other_types[] = { GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, - // GGML_TYPE_TQ1_0, GGML_TYPE_TQ2_0, // TODO: implement for all backends + GGML_TYPE_TQ2_0, + // GGML_TYPE_TQ1_0, // TODO: implement for all backends GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_XS, From a4a4c51f3d40e086b59b73b631b5c43c8fbf4504 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 12 Aug 2026 08:08:19 +0300 Subject: [PATCH 030/211] tests : update speculative params (#26925) --- tools/server/tests/unit/test_speculative.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 0184085b42b2..5837195006bc 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -27,8 +27,8 @@ def test_with_and_without_draft(): global server request = { "prompt": "I believe the meaning of life is", - "temperature": 0.8, - "top_k": 40, + "temperature": 0.2, + "top_k": 5, "seed": 4242, "n_predict": 16, "return_tokens": True, @@ -36,7 +36,6 @@ def test_with_and_without_draft(): server.model_draft = None # disable draft model server.spec_type = None - server.backend_sampling = True server.start() res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 @@ -45,7 +44,6 @@ def test_with_and_without_draft(): # create new server with draft model create_server() - server.backend_sampling = True server.start() res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 From 89e0aa6fd362617d9073e0dafc18e41241521572 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Tue, 11 Aug 2026 23:10:27 -0700 Subject: [PATCH 031/211] opencl: default FA c8 cluster width to 16 on X1E (#26433) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 5d0324dad0ca..257908605994 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -4929,8 +4929,13 @@ static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_ const int x = (e && e[0]) ? atoi(e) : 0; return (x == 8 || x == 16 || x == 32) ? x : 0; // 0 = per-gen default }(); + // X2E needs 16 to keep per-lane o_acc at 128B (the compiler spills the + // kernel-default width); X1E does not spill, but C=16 is still a measured + // +28-30% DK128-GQA4 decode win there (X1-85, kv 4096/8192), neutral on + // DK64 / GQA1 / quant-KV. const int fa_cl_c_gqa4 = fa_cl_c_env ? fa_cl_c_env - : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 16 : 0); + : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E || + backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E ? 16 : 0); const std::string opts_cl_c_gqa4 = fa_cl_c_gqa4 ? " -D FA_CL_C=" + std::to_string(fa_cl_c_gqa4) : std::string(); const std::string fa_cl_c_g8_val = std::to_string(fa_cl_c_gqa4 ? fa_cl_c_gqa4 * 2 : 16); From 4dd127584b87d1b12b6a33c2213197234503d8f7 Mon Sep 17 00:00:00 2001 From: parabelboi Date: Wed, 12 Aug 2026 12:03:32 +0200 Subject: [PATCH 032/211] ui: add read_media tool (#25877) * server: add read_image tool (#25875) Adds a server-tool that allows vision models to analyze server-side images. This tool is reading a single file for now: The image data is base64 encoded and passed to the UI, which decodes it, fills the tag and removes the data URI before passing the tool result back to the model. * cleanup read_image tool: move magic strings to constants * Add dedicated constants file: tools/ui/src/lib/constants/read-image.ts with PREFIX_IMAGE, PREFIX_SIZE, PREFIX_MIME constants * Use ATTACHMENT_SAVED_REGEX from agentic.ts in ChatMessageToolCallBlockReadImage.svelte * Use NEWLINE constant from code.ts instead of hardcoded '\n' * Use PREFIX_SIZE in regex pattern for size parsing * Add SERVER_TOOL_READ_IMAGE_PREFIX_* constants in C++ server-tools.cpp to match the TypeScript PREFIX_* constants for consistency * server: rename read_image tool to read_media for images and audio * Rename server_tool_read_image to server_tool_read_media in C++ * Rename enum BuiltInTool.READ_IMAGE to READ_MEDIA * Rename UI constants, parser, and Svelte component files * Update display label from 'Read image' to 'Read media' * ui: consolidate audio data URI handling into shared utility * Extract getAudioInputFormat to a shared utility (was duplicated inline) * Store raw base64 in base64Data on the message object * Use base64Data to construct data URIs for audio rendering * Update agentic store to build INPUT_AUDIO parts from base64Data * server: read_media: restrict audio to wav/mp3 and minor fixes * Server get_mime_from_extension now only advertises audio/wav and audio/mpeg (the only formats the model's input_audio API accepts) * Case-insensitive extension matching (fixes .MP3, .Wav, etc.) * Unknown extensions return an error instead of a multi-MB data URI that inflates model context with garbage * Updated tool description to document supported formats * Frontend AUDIO_MIME_TO_EXTENSION trimmed to match server * fix a missing import in tools/ui/src/lib/stores/agentic.svelte.ts * server: read_media: add to --tools help text and README tool list * ui: fix indentation in ChatMessageToolCallBlockDefault.svelte * server: read_media tool: fix a cast to use the correct type * server: read_media: multiple fixes * server-tools.cpp import cctype, remove UTF-8 char, check mime before reading file * ui: add MimeTypePrefix.AUDIO and use it in agentic.svelte.ts * server: make read_media inherit from read_file and add uses_cwd * ui: fix formating issues * rm from server * move it to frontend-only tool * correct partial commit * rm unused * ui: address review from allozaur Replace the magic strings, regexes and number in the read_media parser and service with named constants. Path splitting reuses FILE_PATH_SEPARATOR_REGEX, the size header regex moves to READ_MEDIA_SIZE_REGEX derived from PREFIX_SIZE, and FILE_EXTENSION_SEPARATOR lands next to it in constants/code.ts. --------- Co-authored-by: ckrafft Co-authored-by: Xuan Son Nguyen Co-authored-by: Pascal --- tools/server/server-tools.cpp | 30 +++++ .../ChatMessageToolCallBlock.svelte | 3 + .../ChatMessageToolCallBlockDefault.svelte | 37 ++++-- ...essageToolCallBlockExecShellCommand.svelte | 10 +- .../ChatMessageToolCallBlockReadMedia.svelte | 99 ++++++++++++++++ .../ChatMessageToolCall/parsers/read-media.ts | 56 +++++++++ tools/ui/src/lib/constants/built-in-tools.ts | 2 + tools/ui/src/lib/constants/code.ts | 3 + tools/ui/src/lib/constants/index.ts | 1 + tools/ui/src/lib/constants/mcp-resource.ts | 20 +++- tools/ui/src/lib/constants/read-media.ts | 66 +++++++++++ tools/ui/src/lib/constants/tools.ts | 6 + tools/ui/src/lib/enums/files.enums.ts | 1 + tools/ui/src/lib/enums/tools.enums.ts | 1 + tools/ui/src/lib/services/chat.service.ts | 26 +--- .../ui/src/lib/services/read-media.service.ts | 112 ++++++++++++++++++ tools/ui/src/lib/services/tools.service.ts | 16 ++- tools/ui/src/lib/stores/agentic.svelte.ts | 49 +++++++- tools/ui/src/lib/stores/models.svelte.ts | 5 +- tools/ui/src/lib/stores/tools.svelte.ts | 41 ++++++- tools/ui/src/lib/utils/agentic.ts | 19 +-- tools/ui/src/lib/utils/audio-format.ts | 22 ++++ tools/ui/src/lib/utils/index.ts | 5 +- tools/ui/tests/client/README-perf.md | 12 +- tools/ui/tests/unit/agentic-hotpath.bench.ts | 10 +- 25 files changed, 579 insertions(+), 73 deletions(-) create mode 100644 tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockReadMedia.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-media.ts create mode 100644 tools/ui/src/lib/constants/read-media.ts create mode 100644 tools/ui/src/lib/services/read-media.service.ts create mode 100644 tools/ui/src/lib/utils/audio-format.ts diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index a4c1059fff49..fd0ff8ddd7cf 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -1,6 +1,7 @@ #include "server-tools.h" #include "subproc.h" +#include "base64.hpp" #include #include @@ -864,6 +865,7 @@ static bool path_glob_match(const std::string & pattern, const std::string & rel // static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE = 16 * 1024; // 16 KB +static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64 = 32 * 1024 * 1024; // 32 MB struct server_tool_read_file : server_tool { server_tool_read_file() { @@ -899,6 +901,8 @@ struct server_tool_read_file : server_tool { int start_line = json_value(params, "start_line", 1); int end_line = json_value(params, "end_line", -1); // -1 = no limit bool append_loc = json_value(params, "append_loc", false); + // comes from the x-resp-type header, the model cannot ask for it + bool as_base64 = json_value(params, "resp_type", std::string()) == "base64"; auto io = make_tools_io(params); @@ -906,6 +910,23 @@ struct server_tool_read_file : server_tool { if (!io->file_size(path, file_size)) { return {{"error", "cannot stat file: " + path}}; } + + if (as_base64) { + if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64) { + return {{"error", string_format( + "file too large (%zu bytes, max %zu)", + (size_t)file_size, SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64)}}; + } + std::string content; + if (!io->read_file(path, content)) { + return {{"error", "failed to open file: " + path}}; + } + return { + {"base64", base64::encode(content.data(), content.size())}, + {"size_bytes", (size_t) content.size()}, + }; + } + if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE && end_line == -1) { return {{"error", string_format( "file too large (%zu bytes, max %zu). Use start_line/end_line to read a portion.", @@ -2135,6 +2156,15 @@ void server_tools::setup(const std::vector & enabled_tools, params["runtime"] = runtime->spec(); } + // x-resp-type header is only used by read_file for now + if (params.contains("resp_type")) { + params.erase("resp_type"); + } + auto resp_type = get_header(req.headers, "x-resp-type"); + if (!resp_type.empty()) { + params["resp_type"] = resp_type; + } + server_tool & tool = find_tool(tools, tool_name, stream); if (stream) { diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte index 5bbaa4ea311b..7a00f1a8eaa6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte @@ -7,6 +7,7 @@ import ChatMessageToolCallBlockGetInfo from './ChatMessageToolCallBlockGetInfo.svelte'; import ChatMessageToolCallBlockGrepSearch from './ChatMessageToolCallBlockGrepSearch.svelte'; import ChatMessageToolCallBlockReadFile from './ChatMessageToolCallBlockReadFile.svelte'; + import ChatMessageToolCallBlockReadMedia from './ChatMessageToolCallBlockReadMedia.svelte'; import ChatMessageToolCallBlockRunJavascript from './ChatMessageToolCallBlockRunJavascript.svelte'; import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte'; import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte'; @@ -45,6 +46,8 @@ {:else if section.toolName === BuiltInTool.READ_FILE} +{:else if section.toolName === BuiltInTool.READ_MEDIA} + {:else if section.toolName === BuiltInTool.EDIT_FILE} {:else if section.toolName === BuiltInTool.WRITE_FILE} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte index 34dfde78b895..87208e2daa61 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockDefault.svelte @@ -8,14 +8,16 @@ import { MarkdownContent, SyntaxHighlightedCode } from '$lib/components/app'; import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants'; import { getBuiltinToolUi } from '$lib/constants/built-in-tools'; - import { FileTypeText, ToolResultKind } from '$lib/enums'; + import { AttachmentType, FileTypeText, MimeTypeAudio, ToolResultKind } from '$lib/enums'; import type { DatabaseMessageExtra } from '$lib/types'; import { type AgenticSection, classifyToolResult, formatJsonPretty, - parseToolResultWithImages + parseToolResultWithMedia, + type ToolResultLine } from '$lib/utils'; + import { createBase64DataUrl } from '$lib/utils/data-url'; interface Props { section: AgenticSection; @@ -29,8 +31,8 @@ const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? ''); const outputKind = $derived(classifyToolResult(section.toolResult)); - const parsedLines = $derived( - section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : [] + const parsedLines: ToolResultLine[] = $derived( + section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : [] ); @@ -103,13 +105,26 @@
{line.text}
- {#if line.image} - {line.image.name} + {#if line.media} + {#if line.media.type === AttachmentType.AUDIO} + {@const audioMimeType = line.media.mimeType ?? MimeTypeAudio.MP3_MPEG} +
+ +
+ {:else} + {line.media.name} + {/if} {/if} {/each} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte index b2fa8b331c3f..907a3cd12ab6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte @@ -23,7 +23,7 @@ isExitCodeSummaryLine, parseExecShellCommandError, parseExecShellCommandExitStatus, - parseToolResultWithImages, + parseToolResultWithMedia, type ToolResultLine } from '$lib/utils'; @@ -53,7 +53,7 @@ ); const parsedLines: ToolResultLine[] = $derived( - section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : [] + section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : [] ); // Drop the trailing "[exit code: N]" line - rendered as a colored @@ -223,10 +223,10 @@ > {#each outputLines as line, i (i)}
{line.text}
- {#if line.image} + {#if line.media} {line.image.name} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockReadMedia.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockReadMedia.svelte new file mode 100644 index 000000000000..424c794d0bae --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockReadMedia.svelte @@ -0,0 +1,99 @@ + + + + {#snippet titleSnippet()} + Read media + {readMediaMeta?.fileName} + {/snippet} + + {#snippet children(_meta, _ctx)} + {#if section.toolResult} + {#if !mediaAttachment} +
+ Media attachment not found in message extras +
+ {:else if mediaAttachment.type === AttachmentType.AUDIO} +
+ +
+ {:else} +
+ {readMediaMeta?.fileName +
+ {/if} + + {#if readMediaMeta?.sizeBytes || readMediaMeta?.mimeType} +
+ {#if readMediaMeta?.sizeBytes} + Size: {readMediaMeta.sizeBytes} bytes + {/if} + {#if readMediaMeta?.mimeType} + MIME: {readMediaMeta.mimeType} + {/if} +
+ {/if} + + {#if readMediaMeta?.path} +
{readMediaMeta.path}
+ {/if} + {:else} +
+ Waiting for media data... +
+ {/if} + {/snippet} +
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-media.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-media.ts new file mode 100644 index 000000000000..ab1a4d6aa01b --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-media.ts @@ -0,0 +1,56 @@ +import { FILE_PATH_SEPARATOR_REGEX, NEWLINE } from '$lib/constants/code'; +import { + PREFIX_FILE, + PREFIX_MIME, + PREFIX_SIZE, + READ_MEDIA_SIZE_REGEX +} from '$lib/constants/read-media'; +import type { AgenticSection } from '$lib/utils'; + +export interface ReadMediaMeta { + fileName: string; + path: string; + sizeBytes?: number; + mimeType?: string; +} + +/** + * Parse read_media tool result to extract metadata. + * Expected format (after extractBase64Attachments processing): + * File: /path/to/file.png + * Size: 12345 bytes + * MIME: image/png + * [Attachment saved: mcp-attachment-xxx.png] + * + * The data URI line is replaced by the attachment marker by + * agenticStore.extractBase64Attachments before storage. + */ +export function parseReadMediaMeta(section: AgenticSection): ReadMediaMeta | null { + if (!section.toolResult) return null; + + const lines = section.toolResult.split(NEWLINE); + + let fileName = ''; + let path = ''; + let sizeBytes: number | undefined; + let mimeType: string | undefined; + + for (const line of lines) { + const trimmed = line.trim(); + + if (trimmed.startsWith(PREFIX_FILE)) { + path = trimmed.slice(PREFIX_FILE.length).trim(); + fileName = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? path; + } else if (trimmed.startsWith(PREFIX_SIZE)) { + const match = trimmed.match(READ_MEDIA_SIZE_REGEX); + + if (match) sizeBytes = Number(match[1]); + } else if (trimmed.startsWith(PREFIX_MIME)) { + mimeType = trimmed.slice(PREFIX_MIME.length).trim(); + } + } + + if (!path) return null; + + return { fileName, mimeType, path, sizeBytes }; +} diff --git a/tools/ui/src/lib/constants/built-in-tools.ts b/tools/ui/src/lib/constants/built-in-tools.ts index 89bcbecb8ae4..5bd24ffe8397 100644 --- a/tools/ui/src/lib/constants/built-in-tools.ts +++ b/tools/ui/src/lib/constants/built-in-tools.ts @@ -10,6 +10,7 @@ import { Braces, Clock, + Eye, FilePen, FilePlus, FileSearch, @@ -47,6 +48,7 @@ export const BUILTIN_TOOL_UI: Readonly> source: ToolSource.BUILTIN }, [BuiltInTool.READ_FILE]: { icon: FileText, label: 'Read file', source: ToolSource.BUILTIN }, + [BuiltInTool.READ_MEDIA]: { icon: Eye, label: 'Read media', source: ToolSource.FRONTEND }, [BuiltInTool.RUN_JAVASCRIPT]: { icon: Braces, label: 'Run JavaScript', diff --git a/tools/ui/src/lib/constants/code.ts b/tools/ui/src/lib/constants/code.ts index e57e1e6ec575..4b41142001fb 100644 --- a/tools/ui/src/lib/constants/code.ts +++ b/tools/ui/src/lib/constants/code.ts @@ -18,6 +18,9 @@ export const TRIM_TRAILING_PADDING_REGEX = /(?:\n[ \t]*)+$/; // `C:\foo\bar.txt`. Used wherever a parameter accepts a user-supplied path. export const FILE_PATH_SEPARATOR_REGEX = /[\\/]/; +// Separates a file name from its extension, e.g. the '.' in `cover.png`. +export const FILE_EXTENSION_SEPARATOR = '.'; + // Matches the `text:` prefix that file-type identifiers use to denote a // plain-text language (e.g. `text:typescript`). Used by tool-call renderers // to recover the underlying highlight.js language. diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index 09006cf1e991..357a33a6241e 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -49,6 +49,7 @@ export * from './sse'; export * from './precision'; export * from './processing-info'; export * from './pwa'; +export * from './read-media'; export * from './routes'; export * from './sandbox'; export * from './settings-keys'; diff --git a/tools/ui/src/lib/constants/mcp-resource.ts b/tools/ui/src/lib/constants/mcp-resource.ts index 0ef8d9ec13b6..c2639daa12d5 100644 --- a/tools/ui/src/lib/constants/mcp-resource.ts +++ b/tools/ui/src/lib/constants/mcp-resource.ts @@ -1,4 +1,4 @@ -import { MimeTypeImage } from '$lib/enums'; +import { MimeTypeAudio, MimeTypeImage } from '$lib/enums'; // File extension patterns for resource type detection export const IMAGE_FILE_EXTENSION_REGEX = /\.(png|jpg|jpeg|gif|svg|webp)$/i; @@ -27,6 +27,9 @@ export const MCP_RESOURCE_ATTACHMENT_ID_PREFIX = 'res'; // Default file extension for unknown image types export const DEFAULT_IMAGE_EXTENSION = 'img'; +// Default file extension for unknown audio types +export const DEFAULT_AUDIO_EXTENSION = 'mp3'; + // Default filename for resource content downloads export const DEFAULT_RESOURCE_FILENAME = 'resource.txt'; @@ -53,3 +56,18 @@ export const IMAGE_MIME_TO_EXTENSION: Record = { [MimeTypeImage.PNG]: 'png', [MimeTypeImage.WEBP]: 'webp' } as const; + +/** + * Mapping from audio MIME types to file extensions. + * Used for generating attachment filenames from MIME types. + */ +export const AUDIO_MIME_TO_EXTENSION: Record = { + [MimeTypeAudio.MP3]: 'mp3', + [MimeTypeAudio.MP3_MPEG]: 'mp3', + [MimeTypeAudio.VND_WAVE]: 'wav', + [MimeTypeAudio.WAV]: 'wav', + [MimeTypeAudio.WAVE]: 'wav', + [MimeTypeAudio.X_PN_WAV]: 'wav', + [MimeTypeAudio.X_WAV]: 'wav', + [MimeTypeAudio.X_WAVE]: 'wav' +} as const; diff --git a/tools/ui/src/lib/constants/read-media.ts b/tools/ui/src/lib/constants/read-media.ts new file mode 100644 index 000000000000..525c5e90292e --- /dev/null +++ b/tools/ui/src/lib/constants/read-media.ts @@ -0,0 +1,66 @@ +import { + BuiltInTool, + JsonSchemaType, + MimeTypeAudio, + MimeTypeImage, + ToolCallType +} from '$lib/enums'; +import type { OpenAIToolDefinition } from '$lib/types'; + +export const READ_MEDIA_TOOL_NAME = BuiltInTool.READ_MEDIA; + +// header lines of the tool result, parsed back by the read_media renderer +export const PREFIX_FILE = 'File: '; +export const PREFIX_SIZE = 'Size: '; +export const PREFIX_MIME = 'MIME: '; + +/** Byte count of the `Size: ` header line, e.g. `Size: 12345 bytes` -> capture group 1 is `12345`. */ +export const READ_MEDIA_SIZE_REGEX = new RegExp(`^${PREFIX_SIZE}\\s*(\\d+)\\s*bytes`); + +/** Image extensions the tool accepts. The server decodes images with stb_image, which has no webp or tiff. */ +export const READ_MEDIA_IMAGE_MIME: Record = { + gif: MimeTypeImage.GIF, + jpeg: MimeTypeImage.JPEG, + jpg: MimeTypeImage.JPEG, + png: MimeTypeImage.PNG +} as const; + +/** Audio extensions the tool accepts. The `input_audio` API only takes wav and mp3. */ +export const READ_MEDIA_AUDIO_MIME: Record = { + mp3: MimeTypeAudio.MP3_MPEG, + wav: MimeTypeAudio.WAV +} as const; + +/** + * Build the read_media tool definition for the modalities the active model has. + * At least one of the two flags must be true, otherwise the tool is not offered + * at all - a model that cannot see or hear has nothing to do with the bytes. + */ +export function buildReadMediaToolDefinition( + supportsVision: boolean, + supportsAudio: boolean +): OpenAIToolDefinition { + const kinds: string[] = []; + + if (supportsVision) kinds.push(`images (${Object.keys(READ_MEDIA_IMAGE_MIME).join(', ')})`); + + if (supportsAudio) kinds.push(`audio (${Object.keys(READ_MEDIA_AUDIO_MIME).join(', ')})`); + + return { + function: { + description: `Read a media file and attach it to the conversation so it can be perceived directly. Supports ${kinds.join(' and ')}.`, + name: READ_MEDIA_TOOL_NAME, + parameters: { + properties: { + path: { + description: 'Path to the media file', + type: JsonSchemaType.STRING + } + }, + required: ['path'], + type: JsonSchemaType.OBJECT + } + }, + type: ToolCallType.FUNCTION + }; +} diff --git a/tools/ui/src/lib/constants/tools.ts b/tools/ui/src/lib/constants/tools.ts index 65f4457c969d..a467cda2b97a 100644 --- a/tools/ui/src/lib/constants/tools.ts +++ b/tools/ui/src/lib/constants/tools.ts @@ -3,6 +3,12 @@ import { ToolSource } from '$lib/enums/tools.enums'; /** HTTP header carrying the working directory a tool call runs in. The server resolves relative paths against it; the model cannot override it. */ export const X_TOOL_CWD_HEADER = 'x-tool-cwd'; +/** HTTP header asking the server to encode a tool's output differently, e.g. read_file returning base64. Not a tool parameter, so it stays out of the definition the model sees. */ +export const X_RESP_TYPE_HEADER = 'x-resp-type'; + +/** `X_RESP_TYPE_HEADER` value that makes read_file return the raw bytes as base64 instead of text. */ +export const RESP_TYPE_BASE64 = 'base64'; + export const TOOL_GROUP_LABELS = { [ToolSource.BUILTIN]: 'Built-in', [ToolSource.CUSTOM]: 'JSON Schema', diff --git a/tools/ui/src/lib/enums/files.enums.ts b/tools/ui/src/lib/enums/files.enums.ts index eecb36c23e63..5785428cf032 100644 --- a/tools/ui/src/lib/enums/files.enums.ts +++ b/tools/ui/src/lib/enums/files.enums.ts @@ -163,6 +163,7 @@ export enum FileExtensionText { // MIME type prefixes and includes for content detection export enum MimeTypePrefix { IMAGE = 'image/', + AUDIO = 'audio/', TEXT = 'text' } diff --git a/tools/ui/src/lib/enums/tools.enums.ts b/tools/ui/src/lib/enums/tools.enums.ts index 7a9751ee7993..31c992fef4f7 100644 --- a/tools/ui/src/lib/enums/tools.enums.ts +++ b/tools/ui/src/lib/enums/tools.enums.ts @@ -37,6 +37,7 @@ export enum GlobSearchType { */ export enum BuiltInTool { READ_FILE = 'read_file', + READ_MEDIA = 'read_media', EDIT_FILE = 'edit_file', WRITE_FILE = 'write_file', GET_DATETIME = 'get_datetime', diff --git a/tools/ui/src/lib/services/chat.service.ts b/tools/ui/src/lib/services/chat.service.ts index 775a5d6e060e..540d742393f9 100644 --- a/tools/ui/src/lib/services/chat.service.ts +++ b/tools/ui/src/lib/services/chat.service.ts @@ -1,4 +1,5 @@ import { settingsStore } from '../stores/settings.svelte'; +import { getAudioInputFormat } from '../utils/audio-format'; import { capImageDataURLSize } from '../utils/cap-img-size'; import { API_CHAT, @@ -20,18 +21,12 @@ import { import { AttachmentType, ContentPartType, - FileTypeAudio, MessageRole, - MimeTypeAudio, ReasoningFormat, StreamConnectionState } from '$lib/enums'; import { modelsStore } from '$lib/stores/models.svelte'; -import type { - AudioInputFormat, - DatabaseMessageExtraMcpPrompt, - DatabaseMessageExtraMcpResource -} from '$lib/types'; +import type { DatabaseMessageExtraMcpPrompt, DatabaseMessageExtraMcpResource } from '$lib/types'; import type { ApiChatCompletionToolCall, ApiChatMessageContentPart, @@ -43,23 +38,6 @@ import { getAuthHeaders, getJsonHeaders } from '$lib/utils/api-headers'; import { formatAttachmentText } from '$lib/utils/formatters'; import { streamIdentity } from '$lib/utils/stream-identity'; -function getAudioInputFormat(mimeType: string): AudioInputFormat { - const normalizedMimeType = mimeType.trim().toLowerCase(); - - if ( - normalizedMimeType === MimeTypeAudio.WAV || - normalizedMimeType === MimeTypeAudio.WAVE || - normalizedMimeType === MimeTypeAudio.X_WAV || - normalizedMimeType === MimeTypeAudio.X_WAVE || - normalizedMimeType === MimeTypeAudio.VND_WAVE || - normalizedMimeType === MimeTypeAudio.X_PN_WAV - ) { - return FileTypeAudio.WAV; - } - - return FileTypeAudio.MP3; -} - interface ResumableStreamState { bytesReceived: number; updatedAt: number; diff --git a/tools/ui/src/lib/services/read-media.service.ts b/tools/ui/src/lib/services/read-media.service.ts new file mode 100644 index 000000000000..fd66350d000b --- /dev/null +++ b/tools/ui/src/lib/services/read-media.service.ts @@ -0,0 +1,112 @@ +import { ToolsService } from './tools.service'; +import { + FILE_EXTENSION_SEPARATOR, + FILE_PATH_SEPARATOR_REGEX, + NEWLINE, + PREFIX_FILE, + PREFIX_MIME, + PREFIX_SIZE, + READ_MEDIA_AUDIO_MIME, + READ_MEDIA_IMAGE_MIME, + RESP_TYPE_BASE64 +} from '$lib/constants'; +import { BuiltInTool, ToolResponseField } from '$lib/enums'; +import type { ToolExecutionResult } from '$lib/types'; + +/** Modalities of the model the tool call runs for. */ +export interface ReadMediaCapabilities { + audio: boolean; + vision: boolean; +} + +/** Lowercase extension of a path, without the dot. Empty when the file name has none. */ +function fileExtension(path: string): string { + const name = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? ''; + const dot = name.lastIndexOf(FILE_EXTENSION_SEPARATOR); + + return dot > 0 ? name.slice(dot + 1).toLowerCase() : ''; +} + +/** + * **ReadMediaService** - frontend executor for the `read_media` tool + * + * The tool is synthetic: no such tool exists on the server. It reads the file + * through the built-in `read_file` tool with the `base64` response type, then + * turns the bytes into a data URI line. The agentic store lifts that line into + * an image or audio attachment on the tool result message, which is what makes + * the model perceive the file instead of reading a wall of base64. + * + * Living in the frontend is what lets it exist only for models that can + * actually use the result - the server has no idea which model is selected. + * + * @see buildReadMediaToolDefinition in constants/read-media.ts - tool schema sent to the LLM + * @see agenticStore in stores/agentic.svelte.ts - tool dispatch and attachment extraction + */ +export class ReadMediaService { + static async executeTool( + params: Record, + capabilities: ReadMediaCapabilities, + signal?: AbortSignal, + cwd?: string + ): Promise { + const path = typeof params.path === 'string' ? params.path : ''; + + if (!path) { + return { content: 'Error: missing "path" argument.', isError: true }; + } + + const extension = fileExtension(path); + const imageMime = READ_MEDIA_IMAGE_MIME[extension]; + const audioMime = READ_MEDIA_AUDIO_MIME[extension]; + + let resolvedMime: string | undefined; + + if (imageMime && capabilities.vision) resolvedMime = imageMime; + else if (audioMime && capabilities.audio) resolvedMime = audioMime; + + if (!resolvedMime) { + const supported = [ + ...(capabilities.vision ? Object.keys(READ_MEDIA_IMAGE_MIME) : []), + ...(capabilities.audio ? Object.keys(READ_MEDIA_AUDIO_MIME) : []) + ]; + // an unreadable-by-this-model file is a dead end, so say why instead of failing silently + const reason = + imageMime || audioMime + ? `the current model cannot perceive ".${extension}" files` + : `".${extension}" is not a supported media type`; + + return { + content: `Error: ${reason}. Supported: ${supported.join(', ')}.`, + isError: true + }; + } + + const raw = await ToolsService.executeToolRaw( + BuiltInTool.READ_FILE, + { path }, + signal, + cwd, + RESP_TYPE_BASE64 + ); + + if (ToolResponseField.ERROR in raw) { + return { content: String(raw[ToolResponseField.ERROR]), isError: true }; + } + + const base64 = typeof raw.base64 === 'string' ? raw.base64 : ''; + + if (!base64) { + return { content: `Error: no data returned for ${path}.`, isError: true }; + } + + const sizeBytes = typeof raw.size_bytes === 'number' ? raw.size_bytes : 0; + const content = [ + `${PREFIX_FILE}${path}`, + `${PREFIX_SIZE}${sizeBytes} bytes`, + `${PREFIX_MIME}${resolvedMime}`, + `data:${resolvedMime};base64,${base64}` + ].join(NEWLINE); + + return { content, isError: false }; + } +} diff --git a/tools/ui/src/lib/services/tools.service.ts b/tools/ui/src/lib/services/tools.service.ts index cd6b12cebe06..9cb5ecb46dbf 100644 --- a/tools/ui/src/lib/services/tools.service.ts +++ b/tools/ui/src/lib/services/tools.service.ts @@ -1,5 +1,5 @@ import { base } from '$app/paths'; -import { API_TOOLS, X_TOOL_CWD_HEADER } from '$lib/constants'; +import { API_TOOLS, X_RESP_TYPE_HEADER, X_TOOL_CWD_HEADER } from '$lib/constants'; import { ToolResponseField } from '$lib/enums'; import type { ServerBuiltinToolInfo, ToolExecutionResult } from '$lib/types'; import { apiFetch } from '$lib/utils'; @@ -51,16 +51,26 @@ export class ToolsService { * Execute a built-in tool and return the raw JSON response. Unlike * executeTool, this preserves structured fields (e.g. file_glob_search's * `entries` and `base`) that the flattened ToolExecutionResult drops. + * + * @param respType - sent as the x-resp-type request header. Only read_file + * honors it, with `base64` to get the raw bytes instead of decoded text. */ static async executeToolRaw( toolName: string, params: Record, signal?: AbortSignal, - cwd?: string + cwd?: string, + respType?: string ): Promise> { + const headers: Record = {}; + + if (cwd) headers[X_TOOL_CWD_HEADER] = cwd; + + if (respType) headers[X_RESP_TYPE_HEADER] = respType; + return apiFetch>(API_TOOLS.EXECUTE, { body: JSON.stringify({ params, tool: toolName }), - headers: cwd ? { [X_TOOL_CWD_HEADER]: cwd } : undefined, + headers: Object.keys(headers).length > 0 ? headers : undefined, method: 'POST', signal }); diff --git a/tools/ui/src/lib/stores/agentic.svelte.ts b/tools/ui/src/lib/stores/agentic.svelte.ts index 080ed7c89ab7..a84fa34dc276 100644 --- a/tools/ui/src/lib/stores/agentic.svelte.ts +++ b/tools/ui/src/lib/stores/agentic.svelte.ts @@ -22,7 +22,9 @@ import { DEFAULT_AGENTIC_CONFIG, NEWLINE } from '$lib/constants'; import { + AUDIO_MIME_TO_EXTENSION, DATA_URI_BASE64_REGEX, + DEFAULT_AUDIO_EXTENSION, DEFAULT_IMAGE_EXTENSION, IMAGE_MIME_TO_EXTENSION, MCP_ATTACHMENT_NAME_PREFIX @@ -36,6 +38,7 @@ import { ToolCallType } from '$lib/enums'; import { ChatService } from '$lib/services'; +import { ReadMediaService } from '$lib/services/read-media.service'; import { SandboxService } from '$lib/services/sandbox.service'; import { ToolsService } from '$lib/services/tools.service'; import { conversationsStore } from '$lib/stores/conversations.svelte'; @@ -75,9 +78,10 @@ import type { import type { DatabaseMessage, DatabaseMessageExtra, + DatabaseMessageExtraAudioFile, DatabaseMessageExtraImageFile } from '$lib/types/database'; -import { isAbortError } from '$lib/utils'; +import { getAudioInputFormat, isAbortError } from '$lib/utils'; import { SvelteMap } from 'svelte/reactivity'; function createDefaultSession(): AgenticSession { @@ -900,7 +904,18 @@ class AgenticStore { if (executionResult.isError) toolSuccess = false; } else if (toolSource === ToolSource.FRONTEND) { const args = this.parseToolArguments(toolCall.function.arguments); - const executionResult = await SandboxService.executeTool(toolName, args, signal); + const executionResult = + toolName === BuiltInTool.READ_MEDIA + ? await ReadMediaService.executeTool( + args, + { + audio: modelsStore.modelSupportsAudio(effectiveModel), + vision: modelsStore.modelSupportsVision(effectiveModel) + }, + signal, + conversationsStore.activeConversation?.cwd + ) + : await SandboxService.executeTool(toolName, args, signal); result = executionResult.content; @@ -990,7 +1005,19 @@ class AgenticStore { ]; for (const attachment of attachments) { - if (attachment.type === AttachmentType.IMAGE) { + if (attachment.type === AttachmentType.AUDIO) { + if (modelsStore.modelSupportsAudio(effectiveModel)) { + contentParts.push({ + input_audio: { + data: (attachment as DatabaseMessageExtraAudioFile).base64Data, + format: getAudioInputFormat( + (attachment as DatabaseMessageExtraAudioFile).mimeType + ) + }, + type: ContentPartType.INPUT_AUDIO + }); + } + } else if (attachment.type === AttachmentType.IMAGE) { if (modelsStore.modelSupportsVision(effectiveModel)) { contentParts.push({ image_url: { @@ -1101,6 +1128,18 @@ class AgenticStore { return `[Attachment saved: ${name}]`; } + if (mimeType.startsWith(MimeTypePrefix.AUDIO)) { + // audio extras hold the bare base64, the input_audio part has no room for a data URI + attachments.push({ + base64Data, + mimeType, + name, + type: AttachmentType.AUDIO + }); + + return `[Attachment saved: ${name}]`; + } + return line; }); @@ -1108,7 +1147,9 @@ class AgenticStore { } private buildAttachmentName(mimeType: string, index: number): string { - const extension = IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION; + const extension = mimeType.startsWith(MimeTypePrefix.AUDIO) + ? (AUDIO_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_AUDIO_EXTENSION) + : (IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION); return `${MCP_ATTACHMENT_NAME_PREFIX}-${Date.now()}-${index}.${extension}`; } diff --git a/tools/ui/src/lib/stores/models.svelte.ts b/tools/ui/src/lib/stores/models.svelte.ts index 145de119ce7f..150fb3800f64 100644 --- a/tools/ui/src/lib/stores/models.svelte.ts +++ b/tools/ui/src/lib/stores/models.svelte.ts @@ -18,7 +18,10 @@ import { ModelsService } from '$lib/services/models.service'; import { PropsService } from '$lib/services/props.service'; import { conversationsStore } from '$lib/stores/conversations.svelte'; import { isRouterMode, serverStore } from '$lib/stores/server.svelte'; -import { getAuthHeaders, TTLCache } from '$lib/utils'; +// deep imports, not the '$lib/utils' barrel: it re-exports modules that reach back +// into the stores, and going through it here would read a half-built module +import { getAuthHeaders } from '$lib/utils/api-headers'; +import { TTLCache } from '$lib/utils/cache-ttl'; import { detectThinkingSupport, detectThinkingSupportWithReason diff --git a/tools/ui/src/lib/stores/tools.svelte.ts b/tools/ui/src/lib/stores/tools.svelte.ts index b1a946391a0b..3984dd2bc8d8 100644 --- a/tools/ui/src/lib/stores/tools.svelte.ts +++ b/tools/ui/src/lib/stores/tools.svelte.ts @@ -1,4 +1,5 @@ import { + buildReadMediaToolDefinition, buildSandboxToolDefinition, DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY, HOME_TILDE, @@ -15,6 +16,7 @@ import { } from '$lib/enums'; import { ToolsService } from '$lib/services/tools.service'; import { mcpStore } from '$lib/stores/mcp.svelte'; +import { modelsStore, selectedModelName } from '$lib/stores/models.svelte'; import { config } from '$lib/stores/settings.svelte'; import type { OpenAIToolDefinition, ToolEntry, ToolGroup } from '$lib/types'; import { SvelteMap, SvelteSet } from 'svelte/reactivity'; @@ -168,9 +170,42 @@ class ToolsStore { } get frontendTools(): OpenAIToolDefinition[] { - return config().jsSandboxEnabled - ? [buildSandboxToolDefinition(!!config().symbolicMathEnabled)] - : []; + const tools: OpenAIToolDefinition[] = []; + + if (config().jsSandboxEnabled) { + tools.push(buildSandboxToolDefinition(!!config().symbolicMathEnabled)); + } + + const readMedia = this.readMediaTool(); + + if (readMedia) tools.push(readMedia); + + return tools; + } + + /** + * `read_media` runs in the frontend on top of the server's `read_file`, so it + * exists only when that tool is served and the active model can perceive the + * bytes. The server cannot make this call - it does not know which model the + * conversation uses. + */ + private readMediaTool(): OpenAIToolDefinition | null { + const hasReadFile = this._builtinTools.some( + (def) => def.function.name === BuiltInTool.READ_FILE + ); + + if (!hasReadFile) return null; + + const model = selectedModelName() ?? modelsStore.models[0]?.model ?? ''; + + if (!model) return null; + + const vision = modelsStore.modelSupportsVision(model); + const audio = modelsStore.modelSupportsAudio(model); + + if (!vision && !audio) return null; + + return buildReadMediaToolDefinition(vision, audio); } get customTools(): OpenAIToolDefinition[] { diff --git a/tools/ui/src/lib/utils/agentic.ts b/tools/ui/src/lib/utils/agentic.ts index 9c92a42d40ea..b2dd2cd9ed10 100644 --- a/tools/ui/src/lib/utils/agentic.ts +++ b/tools/ui/src/lib/utils/agentic.ts @@ -50,11 +50,11 @@ export interface AgenticSection { } /** - * Represents a tool result line that may reference an image attachment + * Represents a tool result line that may reference a media attachment (image or audio) */ export type ToolResultLine = { text: string; - image?: DatabaseMessageExtraImageFile; + media?: DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile; }; /** @@ -301,16 +301,16 @@ export function splitSearchSummaryList( return { lines }; } -/** Bounded cache for parseToolResultWithImages results. */ +/** Bounded cache for parseToolResultWithMedia results. */ const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32; const toolResultLinesCache = new Map(); /** - * Parse tool result text into lines, matching image attachments by name. + * Parse tool result text into lines, matching media attachments (images and audio) by name. * Memoized: called per render during streaming on unchanged tool result * strings with unchanged extras. */ -export function parseToolResultWithImages( +export function parseToolResultWithMedia( toolResult: string, extras?: DatabaseMessageExtra[] ): ToolResultLine[] { @@ -332,12 +332,13 @@ export function parseToolResultWithImages( if (!match || !extras) return { text: line }; const attachmentName = match[1]; - const image = extras.find( - (e): e is DatabaseMessageExtraImageFile => - e.type === AttachmentType.IMAGE && e.name === attachmentName + const media = extras.find( + (e): e is DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile => + (e.type === AttachmentType.IMAGE || e.type === AttachmentType.AUDIO) && + e.name === attachmentName ); - return { image, text: line }; + return { media, text: line }; }); if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) { diff --git a/tools/ui/src/lib/utils/audio-format.ts b/tools/ui/src/lib/utils/audio-format.ts new file mode 100644 index 000000000000..4f597aad341a --- /dev/null +++ b/tools/ui/src/lib/utils/audio-format.ts @@ -0,0 +1,22 @@ +import { FileTypeAudio, MimeTypeAudio } from '$lib/enums'; +import type { AudioInputFormat } from '$lib/types/api'; + +/** + * Map a MIME type to the AudioInputFormat expected by the API. + */ +export function getAudioInputFormat(mimeType: string): AudioInputFormat { + const normalizedMimeType = mimeType.trim().toLowerCase(); + + if ( + normalizedMimeType === MimeTypeAudio.WAV || + normalizedMimeType === MimeTypeAudio.WAVE || + normalizedMimeType === MimeTypeAudio.X_WAV || + normalizedMimeType === MimeTypeAudio.X_WAVE || + normalizedMimeType === MimeTypeAudio.VND_WAVE || + normalizedMimeType === MimeTypeAudio.X_PN_WAV + ) { + return FileTypeAudio.WAV; + } + + return FileTypeAudio.MP3; +} diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts index 5762e8e4bb92..c3585bd5f85d 100644 --- a/tools/ui/src/lib/utils/index.ts +++ b/tools/ui/src/lib/utils/index.ts @@ -248,7 +248,7 @@ export { export { deriveAgenticSections, buildAssistantRawOutput, - parseToolResultWithImages, + parseToolResultWithMedia, splitSearchSummaryList, hasAgenticContent, classifyToolResult, @@ -325,3 +325,6 @@ export { uuid } from './uuid'; // CSS utilities export { remToPx } from './css'; + +// Audio format helper (used by agentic store and chat service) +export { getAudioInputFormat } from './audio-format'; diff --git a/tools/ui/tests/client/README-perf.md b/tools/ui/tests/client/README-perf.md index a4e3e4c6446a..198b21956832 100644 --- a/tools/ui/tests/client/README-perf.md +++ b/tools/ui/tests/client/README-perf.md @@ -33,12 +33,12 @@ npx vitest --project=client --run tests/client/agentic-stream.perf.svelte.test.t The point of the harness is the _scaling curve_, not any single number. -| Knob | Reads on | -| --------------------------- | ---------------------------------------------------------------------------------------------------- | -| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. | -| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithImages`, `classifyToolResult`). | -| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. | -| `openCodeFence` | `hljs.highlightAuto` on partial code. | +| Knob | Reads on | +| --------------------------- | --------------------------------------------------------------------------------------------------- | +| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. | +| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithMedia`, `classifyToolResult`). | +| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. | +| `openCodeFence` | `hljs.highlightAuto` on partial code. | Deliberately no hard assertions: CI timing is noisy and the value here is the before/after delta, not a gate. diff --git a/tools/ui/tests/unit/agentic-hotpath.bench.ts b/tools/ui/tests/unit/agentic-hotpath.bench.ts index 8bdaa095a89d..24dcb48911f8 100644 --- a/tools/ui/tests/unit/agentic-hotpath.bench.ts +++ b/tools/ui/tests/unit/agentic-hotpath.bench.ts @@ -6,7 +6,7 @@ // // Run: npx vitest bench --project=unit tests/unit/agentic-hotpath.bench.ts -import { classifyToolResult, parseToolResultWithImages } from '$lib/utils/agentic'; +import { classifyToolResult, parseToolResultWithMedia } from '$lib/utils/agentic'; import { detectIncompleteCodeBlock, highlightCode } from '$lib/utils/code'; import { computeLineDiff } from '$lib/utils/compute-line-diff'; import { preprocessLaTeX } from '$lib/utils/latex-protection'; @@ -200,17 +200,17 @@ describe('exit-code regex', () => { // --- per-line result parsers ---------------------------------------------- -describe('parseToolResultWithImages', () => { +describe('parseToolResultWithMedia', () => { bench('1KB', () => { - parseToolResultWithImages(SHELL_OUTPUT_1KB, []); + parseToolResultWithMedia(SHELL_OUTPUT_1KB, []); }); bench('200KB', () => { - parseToolResultWithImages(SHELL_OUTPUT_200KB, []); + parseToolResultWithMedia(SHELL_OUTPUT_200KB, []); }); bench('2MB', () => { - parseToolResultWithImages(SHELL_OUTPUT_2MB, []); + parseToolResultWithMedia(SHELL_OUTPUT_2MB, []); }); }); From 5d9e5ac30e469d44c0a5a52556de0ead03aaa5b0 Mon Sep 17 00:00:00 2001 From: Chipmunk <101038159+CHIPMUNK-T0T@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:20:28 +0900 Subject: [PATCH 033/211] server : support slot save/restore with media inputs (#26640) * server : save serialized image chunks at the end of the llama state * server : support multimodal slot state save/restore with packed payload * server : refine image slot state serialization * server : support media slot state and centralize media validation * server : remove unnecessary comment * server : remove defensive media checks and move the chunk type check to validate() --- include/llama.h | 1 + src/llama-context.cpp | 11 + tools/server/server-common.cpp | 185 ++++++++++- tools/server/server-common.h | 10 +- tools/server/server-context.cpp | 81 +++-- tools/server/tests/unit/test_slot_save.py | 363 ++++++++++++++++++++-- tools/server/tests/utils.py | 6 + 7 files changed, 586 insertions(+), 71 deletions(-) diff --git a/include/llama.h b/include/llama.h index bfef0e1d1129..ef278c9c3238 100644 --- a/include/llama.h +++ b/include/llama.h @@ -883,6 +883,7 @@ extern "C" { const llama_token * tokens, size_t n_token_count); + // If tokens_out is NULL, only the token count is reported through n_token_count_out and no state is loaded LLAMA_API size_t llama_state_seq_load_file( struct llama_context * ctx, const char * filepath, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0de3a68d1cb0..aa9fb2c3b481 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3110,6 +3110,17 @@ size_t llama_context::state_seq_load_file(llama_seq_id seq_id, const char * file { const uint32_t n_token_count = file.read_u32(); + if (tokens_out == nullptr) { + const size_t n_token_max = (file.size() - file.tell()) / sizeof(llama_token); + if (n_token_count > n_token_max) { + LLAMA_LOG_ERROR("%s: token count in sequence state file exceeds the file size! %u > %zu\n", __func__, n_token_count, n_token_max); + return 0; + } + + *n_token_count_out = n_token_count; + return file.tell(); + } + if (n_token_count > n_token_capacity) { LLAMA_LOG_ERROR("%s: token count in sequence state file exceeded capacity! %u > %zu\n", __func__, n_token_count, n_token_capacity); return 0; diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c9109fc9626e..148b7ea6d81a 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -13,6 +13,8 @@ #include #include #include +#include +#include json format_error_response(const std::string & message, const enum error_type type) { std::string type_str; @@ -235,6 +237,102 @@ static inline raw_buffer base64_decode(const std::string & encoded_string) { // server_tokens implementation // +namespace { + +constexpr uint32_t SERVER_TOKENS_STATE_VERSION = 1; + +uint32_t server_tokens_state_u32(size_t value) { + if (value > std::numeric_limits::max()) { + throw std::runtime_error("Server tokens state is too large"); + } + return value; +} + +class server_tokens_state_writer { +public: + template + void write(T value) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + const auto * ptr = reinterpret_cast(&value); + data.insert(data.end(), ptr, ptr + sizeof(value)); + } + + template + void write(const std::vector & values) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + write(server_tokens_state_u32(values.size())); + if (values.empty()) { + return; + } + const auto * ptr = reinterpret_cast(values.data()); + data.insert(data.end(), ptr, ptr + values.size() * sizeof(T)); + } + + void write_media_chunk(const mtmd_input_chunk * chunk) { + size_t chunk_size = 0; + if (mtmd_input_chunk_save(chunk, nullptr, 0, &chunk_size) != 0 || chunk_size == 0) { + throw std::runtime_error("Cannot serialize media chunk in server tokens"); + } + std::vector chunk_data(server_tokens_state_u32(chunk_size)); + if (mtmd_input_chunk_save(chunk, chunk_data.data(), chunk_data.size(), nullptr) != 0) { + throw std::runtime_error("Cannot serialize media chunk in server tokens"); + } + write(chunk_data); + } + + std::vector take() { + data.resize((data.size() + sizeof(llama_token) - 1) / sizeof(llama_token) * sizeof(llama_token), 0); + return std::move(data); + } + +private: + std::vector data; +}; + +class server_tokens_state_reader { +public: + server_tokens_state_reader(const char * data, size_t size) : data(data), size(size) {} + + template + T read() { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + if (size - pos < sizeof(T)) { + throw std::runtime_error("Unexpected end of server tokens state"); + } + T value; + std::memcpy(&value, data + pos, sizeof(value)); + pos += sizeof(value); + return value; + } + + template + std::vector read_vector() { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + const uint32_t n_values = read(); + // reject before resizing, so that a small corrupted payload cannot request a huge allocation + if (n_values > remaining() / sizeof(T)) { + throw std::runtime_error("Unexpected end of server tokens state"); + } + std::vector values(n_values); + if (n_values > 0) { + std::memcpy(values.data(), data + pos, values.size() * sizeof(T)); + pos += values.size() * sizeof(T); + } + return values; + } + + size_t remaining() const { + return size - pos; + } + +private: + const char * data; + size_t size; + size_t pos = 0; +}; + +} // namespace + server_tokens::server_tokens(mtmd::input_chunks & mtmd_chunks, bool has_mtmd) : has_mtmd(has_mtmd) { for (size_t i = 0; i < mtmd_chunks.size(); ++i) { push_back(mtmd_chunks[i]); @@ -408,6 +506,73 @@ const llama_tokens & server_tokens::get_tokens() const { return tokens; } +std::vector server_tokens::serialize() const { + static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size"); + + server_tokens_state_writer writer; + writer.write((llama_token) LLAMA_TOKEN_NULL); + writer.write(SERVER_TOKENS_STATE_VERSION); + writer.write(tokens); + + std::vector media_keys; + media_keys.reserve(map_idx_to_media.size()); + for (const auto & item : map_idx_to_media) { + media_keys.push_back(server_tokens_state_u32(item.first)); + } + writer.write(media_keys); + + for (const auto & item : map_idx_to_media) { + writer.write_media_chunk(item.second.get()); + } + + return writer.take(); +} + +server_tokens server_tokens::deserialize(const llama_tokens & packed, bool has_mtmd) { + static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size"); + + if (packed.empty() || packed[0] != LLAMA_TOKEN_NULL) { + // plain token list, as written by older versions + return server_tokens(packed, has_mtmd); + } + + server_tokens_state_reader reader(reinterpret_cast(packed.data()), packed.size() * sizeof(llama_token)); + reader.read(); // format marker + if (reader.read() != SERVER_TOKENS_STATE_VERSION) { + throw std::runtime_error("Unsupported server tokens state version"); + } + + const llama_tokens tokens = reader.read_vector(); + + // the media start indices, followed by the media chunks in the same order + const std::vector media_keys = reader.read_vector(); + if (!media_keys.empty() && !has_mtmd) { + throw std::runtime_error("Cannot restore media tokens without an mmproj"); + } + + server_tokens result(tokens, has_mtmd); + + for (const uint32_t key : media_keys) { + const size_t start_idx = key; + const std::vector chunk_data = reader.read_vector(); + if (chunk_data.empty()) { + throw std::runtime_error("Cannot load media chunk from server tokens state"); + } + + mtmd::input_chunk_ptr chunk(mtmd_input_chunk_load(chunk_data.data(), chunk_data.size())); + if (!chunk) { + throw std::runtime_error("Cannot load media chunk from server tokens state"); + } + result.map_idx_to_media[start_idx] = std::move(chunk); + } + + if (reader.remaining() >= sizeof(llama_token)) { + throw std::runtime_error("Trailing data in server tokens state"); + } + + return result; +} + llama_tokens server_tokens::get_text_tokens() const { llama_tokens res; res.reserve(tokens.size()); @@ -530,14 +695,28 @@ bool server_tokens::validate(const struct llama_context * ctx) const { const llama_model * model = llama_get_model(ctx); const llama_vocab * vocab = llama_model_get_vocab(model); const int32_t n_vocab = llama_vocab_n_tokens(vocab); + size_t n_media = 0; for (size_t i = 0; i < tokens.size(); ++i) { const auto & t = tokens[i]; if (t == LLAMA_TOKEN_NULL) { try { const auto & chunk = find_chunk(i); - size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get()); - i += n_tokens - 1; // will be +1 by the for loop + if (mtmd_input_chunk_get_type(chunk.get()) == MTMD_INPUT_CHUNK_TYPE_TEXT) { + return false; + } + const size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get()); + const llama_pos n_pos = mtmd_input_chunk_get_n_pos(chunk.get()); + if (n_tokens == 0 || n_pos <= 0 || n_tokens > tokens.size() - i) { + return false; + } + for (size_t j = i; j < i + n_tokens; ++j) { + if (tokens[j] != LLAMA_TOKEN_NULL) { + return false; + } + } + ++n_media; + i += n_tokens - 1; } catch (const std::exception & e) { return false; } @@ -545,7 +724,7 @@ bool server_tokens::validate(const struct llama_context * ctx) const { return false; } } - return true; + return n_media == map_idx_to_media.size(); } server_tokens server_tokens::clone() const { diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6ef797ebb473..79607e288c0e 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -201,11 +201,14 @@ struct server_tokens { // for compatibility with context shift and prompt truncation void insert(const llama_tokens & inp_tokens); - // for compatibility with speculative decoding, ctx shift, slot save/load + // for compatibility with speculative decoding, ctx shift const llama_tokens & get_tokens() const; llama_tokens get_text_tokens() const; + std::vector serialize() const; + static server_tokens deserialize(const llama_tokens & packed, bool has_mtmd); + // for compatibility with speculative decoding void set_token(llama_pos pos, llama_token id); @@ -213,9 +216,6 @@ struct server_tokens { bool empty() const { return tokens.empty(); } - // true if the sequence actually contains image/audio chunks. - bool has_media() const { return !map_idx_to_media.empty(); } - void clear() { map_idx_to_media.clear(); tokens.clear(); @@ -230,7 +230,7 @@ struct server_tokens { // split the tokens into message spans, skipping over media chunks common_chat_msg_spans find_message_spans(const common_chat_msg_delimiters & delims) const; - // make sure all text tokens are within the vocab range + // check text token IDs and the mapping between media chunks and token ranges bool validate(const struct llama_context * ctx) const; server_tokens clone() const; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d75c6856dfa9..8b2ae72a4cdc 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2072,18 +2072,6 @@ struct server_context_impl { queue_results.send(std::move(res)); } - // Gate slot save/restore/erase on slot content (does it hold media), - // not model capability: a multimodal model may hold a pure-text slot. - bool check_slot_no_media(const server_slot & slot, const int id_task) { - if (slot.prompt.tokens.has_media()) { - send_error(id_task, - "This operation is not supported while the slot holds image/audio tokens (a pure-text prefix is supported)", - ERROR_TYPE_NOT_SUPPORTED); - return false; - } - return true; - } - void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) { auto res = std::make_unique(); @@ -2577,9 +2565,6 @@ struct server_context_impl { send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); break; } - if (!check_slot_no_media(*slot, task.id)) { - break; - } if (slot->is_processing()) { // if requested slot is unavailable, we defer this task for processing later SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); @@ -2592,9 +2577,22 @@ struct server_context_impl { std::string filename = task.slot_action.filename; std::string filepath = task.slot_action.filepath; - const llama_tokens tokens = slot->prompt.tokens.get_text_tokens(); - const size_t token_count = tokens.size(); - const size_t nwrite = llama_state_seq_save_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), token_count); + std::vector packed; + try { + packed = slot->prompt.tokens.serialize(); + } catch (const std::exception & err) { + send_error(task, err.what(), ERROR_TYPE_NOT_SUPPORTED); + break; + } + + GGML_ASSERT(packed.size() % sizeof(llama_token) == 0); + const size_t nwrite = llama_state_seq_save_file( + ctx_tgt, filepath.c_str(), slot->id, + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + if (nwrite == 0) { + send_error(task, "Unable to save slot", ERROR_TYPE_SERVER); + break; + } const int64_t t_end = ggml_time_us(); const double t_save_ms = (t_end - t_start) / 1000.0; @@ -2604,7 +2602,7 @@ struct server_context_impl { res->id_slot = id_slot; res->filename = filename; res->is_save = true; - res->n_tokens = token_count; + res->n_tokens = slot->prompt.tokens.size(); res->n_bytes = nwrite; res->t_ms = t_save_ms; queue_results.send(std::move(res)); @@ -2629,18 +2627,37 @@ struct server_context_impl { std::string filename = task.slot_action.filename; std::string filepath = task.slot_action.filepath; - llama_tokens tokens; - tokens.resize(slot->n_ctx); - size_t token_count = 0; - size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count); - if (nread == 0) { - slot->prompt.clear(); // KV may already been invalidated? - send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST); + size_t nread = 0; + try { + size_t n_packed = 0; + llama_tokens packed; + nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, nullptr, 0, &n_packed); + if (nread != 0) { + packed.resize(std::max(1, n_packed)); + nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, packed.data(), packed.size(), &n_packed); + } + if (nread == 0) { + throw std::runtime_error("No available space in KV cache or invalid slot save file"); + } + packed.resize(n_packed); + + server_tokens restored = server_tokens::deserialize(packed, mctx != nullptr); + + if (restored.size() > (size_t) slot->n_ctx) { + throw std::runtime_error("Restored prompt does not fit in the slot context"); + } + + if (!restored.validate(ctx_tgt)) { + throw std::runtime_error("Invalid tokens in slot save file"); + } + + slot->prompt.clear(); + slot->prompt.tokens = std::move(restored); + } catch (const std::exception & err) { + slot->prompt_clear(); + send_error(task, std::string("Unable to restore slot: ") + err.what(), ERROR_TYPE_INVALID_REQUEST); break; } - tokens.resize(token_count); - slot->prompt.clear(); - slot->prompt.tokens.insert(tokens); const int64_t t_end = ggml_time_us(); const double t_restore_ms = (t_end - t_start) / 1000.0; @@ -2650,7 +2667,7 @@ struct server_context_impl { res->id_slot = id_slot; res->filename = filename; res->is_save = false; - res->n_tokens = token_count; + res->n_tokens = slot->prompt.tokens.size(); res->n_bytes = nread; res->t_ms = t_restore_ms; queue_results.send(std::move(res)); @@ -2663,10 +2680,6 @@ struct server_context_impl { send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); break; } - // Gate on slot content, consistent with save/restore. - if (!check_slot_no_media(*slot, task.id)) { - break; - } if (slot->is_processing()) { // if requested slot is unavailable, we defer this task for processing later SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index be22d9859efd..05acb1be1432 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/tools/server/tests/unit/test_slot_save.py @@ -2,6 +2,10 @@ from utils import * import base64 import requests +import struct + +# sequence state file: magic(4) version(4) payload_size(4), then payload_size llama_token words +STATE_FILE_HEADER_SIZE = 12 server = ServerPreset.tinyllama2() @@ -72,6 +76,60 @@ def test_slot_save_restore(): assert res.body["timings"]["prompt_n"] == 1 +def test_slot_restore_legacy_token_list(): + global server + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "What is the capital of France?", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "slot_legacy.bin", + }) + assert res.status_code == 200 + assert res.body["n_saved"] == 84 + + # rewrite the token payload into a plain token list, as written by servers that predate the packed server_tokens format + path = os.path.join("tmp", "slot_legacy.bin") + with open(path, "rb") as f: + data = bytearray(f.read()) + + # the payload written by this server starts with a packed header: LLAMA_TOKEN_NULL(4) version(4) n_tokens(4) + packed_header_size = 12 + + payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0] + payload_end = STATE_FILE_HEADER_SIZE + payload_size * 4 + n_tokens = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE + 8)[0] + assert n_tokens == 84 + + tokens_start = STATE_FILE_HEADER_SIZE + packed_header_size + data = data[:STATE_FILE_HEADER_SIZE] + data[tokens_start:tokens_start + n_tokens * 4] + data[payload_end:] + struct.pack_into("=I", data, STATE_FILE_HEADER_SIZE - 4, n_tokens) + + with open(path, "wb") as f: + f.write(data) + + # the plain token list must restore, and the restored KV must be reusable + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "slot_legacy.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == 84 + + res = server.make_request("POST", "/completion", data={ + "prompt": "What is the capital of Germany?", + "id_slot": 0, + "cache_prompt": True, + }) + assert res.status_code == 200 + assert res.body["timings"]["prompt_n"] == 6 # only the different part is processed + + + def test_slot_erase(): global server server.start() @@ -103,14 +161,12 @@ def test_slot_erase(): # # Multimodal server (mmproj loaded) slot save/restore. # -# Regression coverage for issue #21133: slot save/restore/erase must be gated on -# the slot's CONTENT (does it actually hold image/audio tokens) rather than the -# model's CAPABILITY (is an mmproj loaded). A pure-text slot on a multimodal -# server must save/restore/erase normally; a slot that actually holds an image -# must be rejected with ERROR_TYPE_NOT_SUPPORTED (HTTP 501). +# A pure-text slot on a multimodal server and a slot containing images must both support save/restore. +# Erase remains gated on the slot's content. # IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png" +IMG_URL_TRUCK = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png" def _get_img_base64(url: str) -> str: @@ -121,8 +177,7 @@ def _get_img_base64(url: str) -> str: @pytest.fixture def mmproj_server(): - # tinygemma3 is a small multimodal model: the mmproj is provided by the HF - # registry API and auto-downloaded on first run. + # tinygemma3 is a small multimodal model: the mmproj is provided by the HF registry API and auto-downloaded on first run. os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>' mm_server = ServerPreset.tinygemma3() mm_server.slot_save_path = "./tmp" @@ -159,10 +214,7 @@ def test_slot_save_restore_text_only_on_multimodal(mmproj_server): assert res.status_code == 200 assert res.body["n_restored"] == n_saved - # The restored slot is usable for a follow-up completion. We do NOT assert - # prefix reuse here: tinygemma3 is a SWA model, which forces full prompt - # re-processing after a restore (a model property, not the save/restore gate - # under test). + # Prefix reuse is not checked with the default SWA cache. res = server.make_request("POST", "/completion", data={ "prompt": "The quick brown fox jumps over the lazy dog.", "id_slot": 0, @@ -171,54 +223,307 @@ def test_slot_save_restore_text_only_on_multimodal(mmproj_server): assert res.status_code == 200 -def test_slot_save_rejected_when_slot_holds_image(mmproj_server): +def test_slot_save_restore_with_image(mmproj_server): server = mmproj_server + # Use the full SWA cache so the restored image prefix can be reused. + server.swa_full = True server.start() - # Process a prompt that actually contains an image on slot 1. + prompt_cat = { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } res = server.make_request("POST", "/completions", data={ "temperature": 0.0, "top_k": 1, "id_slot": 1, "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + content_cat = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + assert res.body["timings"]["cache_n"] == 0 + assert prompt_n_full > 32 # text plus image tokens are all processed + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "mm_slot_image.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + n_written = res.body["n_written"] + assert n_saved > 0 + assert n_written > 0 + + res = server.make_request("POST", "/slots/1?action=erase") + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_image.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + assert res.body["n_read"] == n_written + + # a different image must not reuse the restored image tokens; only the text prefix before the image is common + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, "prompt": { "prompt_string": "What is this: <__media__>\n", - "multimodal_data": [ _get_img_base64(IMG_URL_CAT) ], + "multimodal_data": [_get_img_base64(IMG_URL_TRUCK)], }, }) assert res.status_code == 200 + cache_n = res.body["timings"]["cache_n"] + assert cache_n < 16 + assert res.body["timings"]["prompt_n"] == prompt_n_full - cache_n - # Saving a slot that holds image tokens must be rejected (HTTP 501, - # not_supported_error). - res = server.make_request("POST", "/slots/1?action=save", data={ + # restore again and resend the same image: the image tokens must be reused and greedy sampling must reproduce the original content + res = server.make_request("POST", "/slots/0?action=restore", data={ "filename": "mm_slot_image.bin", }) - assert res.status_code != 200 - assert res.body["error"]["type"] == "not_supported_error" + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content_cat -def test_slot_erase_text_only_on_multimodal(mmproj_server): +def test_slot_save_restore_with_two_images(mmproj_server): server = mmproj_server + server.swa_full = True + server.n_ctx = 2048 # two images need more than the default 512 per slot server.start() - res = server.make_request("POST", "/completion", data={ - "prompt": "The quick brown fox jumps over the lazy dog.", + prompt = { + "prompt_string": "A: <__media__> B: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT), _get_img_base64(IMG_URL_TRUCK)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, "id_slot": 1, "cache_prompt": True, + "prompt": prompt, }) assert res.status_code == 200 - prompt_n = res.body["timings"]["prompt_n"] - assert prompt_n > 0 # all tokens are processed + content = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + assert prompt_n_full > 64 - # Erasing a pure-text slot must succeed even though an mmproj is loaded. - res = server.make_request("POST", "/slots/1?action=erase") + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "mm_slot_two_images.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_two_images.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content + + +def test_slot_save_restore_with_image_across_restart(mmproj_server): + server = mmproj_server + server.swa_full = True + server.start() + + prompt_cat = { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + content = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_restart.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + + # restart the server with the same model and mmproj: the saved file must restore in the new process and the image KV must be reused + server.stop() + server.start() + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_restart.bin", + }) assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content + + +def test_slot_save_restore_image_payload_larger_than_context(mmproj_server): + server = mmproj_server + server.swa_full = True + server.start() + + # the slot context, as the server computed it (n_ctx split across the slots) + res = server.make_request("GET", "/props") + assert res.status_code == 200 + n_ctx_slot = res.body["default_generation_settings"]["n_ctx"] + + # a filler token, used to grow the prompt up to the slot context + res = server.make_request("POST", "/tokenize", data={"content": " hello" * 8}) + assert res.status_code == 200 + assert len(res.body["tokens"]) == 8 + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + }, + }) + assert res.status_code == 200 + + prompt_cat = { + "prompt_string": "What is this: <__media__>\n" + " hello" * (n_ctx_slot - res.body["timings"]["prompt_n"] - 8), + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + prompt_n_full = res.body["timings"]["cache_n"] + res.body["timings"]["prompt_n"] + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_large_payload.bin", + }) + assert res.status_code == 200 + + path = os.path.join("tmp", "mm_slot_large_payload.bin") + with open(path, "rb") as f: + data = bytearray(f.read()) + payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0] + assert payload_size > n_ctx_slot # the scenario under test: the payload does not fit in n_ctx - # Re-running the same prompt should process all tokens again. + # drop the image from the slot, then restore it from the file res = server.make_request("POST", "/completion", data={ - "prompt": "The quick brown fox jumps over the lazy dog.", + "prompt": "The quick brown fox", + "id_slot": 0, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_large_payload.bin", + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + + +def test_slot_restore_media_file_without_mmproj(mmproj_server): + server = mmproj_server + server.start() + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + }, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_no_mmproj.bin", + }) + assert res.status_code == 200 + + # restart the same model without the mmproj: restoring the media file must fail gracefully and leave the slot usable + server.stop() + server.no_mmproj = True + server.start() + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_no_mmproj.bin", + }) + assert res.status_code == 400 + assert "Cannot restore media tokens without an mmproj" in res.body["error"]["message"] + + # A failed restore must leave the slot empty and usable. + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, "id_slot": 1, "cache_prompt": True, + "prompt": "The quick brown fox", + }) + assert res.status_code == 200 + content = res.body["content"] + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": "The quick brown fox", }) assert res.status_code == 200 - assert res.body["timings"]["prompt_n"] == prompt_n # all tokens are processed again + assert res.body["timings"]["cache_n"] == 0 + assert res.body["content"] == content diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index fffe07a67116..9171dbc02977 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -86,6 +86,7 @@ class ServerProcess: server_reranking: bool | None = False server_metrics: bool | None = False kv_unified: bool | None = False + swa_full: bool | None = False server_slots: bool | None = False pooling: str | None = None api_key: str | None = None @@ -106,6 +107,7 @@ class ServerProcess: chat_template_file: str | None = None server_path: str | None = None mmproj_url: str | None = None + no_mmproj: bool | None = None media_path: str | None = None sleep_idle_seconds: int | None = None cache_ram: int | None = None @@ -198,6 +200,8 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.append("--metrics") if self.kv_unified: server_args.append("--kv-unified") + if self.swa_full: + server_args.append("--swa-full") if self.server_slots: server_args.append("--slots") else: @@ -259,6 +263,8 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--chat-template-file", self.chat_template_file]) if self.mmproj_url: server_args.extend(["--mmproj-url", self.mmproj_url]) + if self.no_mmproj: + server_args.append("--no-mmproj") if self.media_path: server_args.extend(["--media-path", self.media_path]) if self.sleep_idle_seconds is not None: From 13fd0bb55ed30216ec47ee615c0a28455baedbf0 Mon Sep 17 00:00:00 2001 From: Daniel Bevenius Date: Wed, 12 Aug 2026 12:46:06 +0200 Subject: [PATCH 034/211] cmake : add config version support (ggml/1582) * cmake : add config version support (wip) [no ci] This commit adds support for find_package using a version, for example: ``` find_package(ggml 0.19.0 REQUIRED) ``` examples/test-cmake has been updated to use this and build scripts have been added to verify this manually. This is still a work in progress and I'm not sure about the scripts and if we can find better ways to test this but it might be useful to have for verification of changes to the cmake build. * cmake : add semver to ggml backends [no ci] This commit adds a semver to the ggml backend modules files. The motivation for this is that the backends are currently loaded just a file extension, for example .so on linux. With the introduction of semantic versioning installing a new version should just work but since these files don't have a version they would get overwritten. Adding the semver to the library names allows multiple version to be supported and the correct one will be loaded by the code. I've only tested this on linux and need to test on mac and win. * Revert "cmake : add semver to ggml backends [no ci]" This reverts commit 53a6c58a07591951324c891b9986b2cffe5c7972. * examples : update build-install.sh and set GGML_BACKEND_DIR --- ggml/CMakeLists.txt | 4 ++-- ggml/cmake/ggml-config.cmake.in | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 1b1de6b7d451..b6843578b704 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -402,7 +402,7 @@ configure_package_config_file( GGML_BIN_INSTALL_DIR) write_basic_package_version_file( - ${CMAKE_CURRENT_BINARY_DIR}/ggml-version.cmake + ${CMAKE_CURRENT_BINARY_DIR}/ggml-config-version.cmake VERSION ${GGML_INSTALL_VERSION} COMPATIBILITY SameMajorVersion) @@ -414,7 +414,7 @@ message(STATUS "ggml version: ${GGML_INSTALL_VERSION}") message(STATUS "ggml commit: ${GGML_BUILD_COMMIT}") install(FILES ${CMAKE_CURRENT_BINARY_DIR}/ggml-config.cmake - ${CMAKE_CURRENT_BINARY_DIR}/ggml-version.cmake + ${CMAKE_CURRENT_BINARY_DIR}/ggml-config-version.cmake DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ggml) if (MSVC) diff --git a/ggml/cmake/ggml-config.cmake.in b/ggml/cmake/ggml-config.cmake.in index 23a3066f56dd..abe17804a5a7 100644 --- a/ggml/cmake/ggml-config.cmake.in +++ b/ggml/cmake/ggml-config.cmake.in @@ -113,6 +113,7 @@ set_and_check(GGML_LIB_DIR "@PACKAGE_GGML_LIB_INSTALL_DIR@") if(NOT TARGET ggml::ggml) find_package(Threads REQUIRED) + unset(GGML_LIBRARY CACHE) find_library(GGML_LIBRARY ggml REQUIRED HINTS ${GGML_LIB_DIR} @@ -121,8 +122,10 @@ if(NOT TARGET ggml::ggml) add_library(ggml::ggml UNKNOWN IMPORTED) set_target_properties(ggml::ggml PROPERTIES - IMPORTED_LOCATION "${GGML_LIBRARY}") + IMPORTED_LOCATION "${GGML_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${GGML_INCLUDE_DIR}") + unset(GGML_BASE_LIBRARY CACHE) find_library(GGML_BASE_LIBRARY ggml-base REQUIRED HINTS ${GGML_LIB_DIR} @@ -132,6 +135,7 @@ if(NOT TARGET ggml::ggml) set_target_properties(ggml::ggml-base PROPERTIES IMPORTED_LOCATION "${GGML_BASE_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${GGML_INCLUDE_DIR}" INTERFACE_LINK_LIBRARIES "${GGML_BASE_INTERFACE_LINK_LIBRARIES}") set(_ggml_all_targets "") @@ -140,6 +144,7 @@ if(NOT TARGET ggml::ggml) string(REPLACE "-" "_" _ggml_backend_pfx "${_ggml_backend}") string(TOUPPER "${_ggml_backend_pfx}" _ggml_backend_pfx) + unset(${_ggml_backend_pfx}_LIBRARY CACHE) find_library(${_ggml_backend_pfx}_LIBRARY ${_ggml_backend} REQUIRED HINTS ${GGML_LIB_DIR} From af05a42a7c18db1f64be406505e394f61589cf56 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 12 Aug 2026 13:59:44 +0300 Subject: [PATCH 035/211] sync : ggml --- scripts/sync-ggml.last | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/sync-ggml.last b/scripts/sync-ggml.last index 7af9ecb0f60a..c08cb625bb41 100644 --- a/scripts/sync-ggml.last +++ b/scripts/sync-ggml.last @@ -1 +1 @@ -30bf8685ed4eb0a47f2b06229543327749904150 +8846b79e66747bb9f68597420e95114c177315ce From ece98b87f72bc9e3bd356a217feda8dd5f9ca24c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Wed, 12 Aug 2026 13:24:10 +0200 Subject: [PATCH 036/211] model : disallow integer dflash sliding_window_pattern (#26900) * fix sliding_window_pattern * disallow integer pattern --- src/models/dflash.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index bfbdb28ee5e9..eb3633676210 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -66,7 +66,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { // DFlash has a single rope, so the SWA rope == main rope. if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; - ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; } From 132753bf4e52b1a8bda8f6ec33f1785bd80470da Mon Sep 17 00:00:00 2001 From: Jonathan Clohessy Date: Wed, 12 Aug 2026 12:49:11 +0100 Subject: [PATCH 037/211] kleidiai: Add runtime feature detection mechanism for aarch64/kleidiai (#26076) * Add runtime feature detection mechanism for aarch64/kleidiai Signed-off-by: Jonathan Clohessy * Address Review Comments Signed-off-by: Jonathan Clohessy * Add log warning for NSMC reserved value Signed-off-by: Jonathan Clohessy * Address review comments Signed-off-by: Jonathan Clohessy * Fix Rebase, move code from cpu-feats to ggml-feats Signed-off-by: Jonathan Clohessy * Address naming of runtime feature struct Signed-off-by: Jonathan Clohessy --------- Signed-off-by: Jonathan Clohessy --- ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp | 87 +------ ggml/src/ggml-cpu/kleidiai/kleidiai.cpp | 316 +++++++++++++++-------- ggml/src/ggml-feats.h | 166 ++++++++++++ 3 files changed, 385 insertions(+), 184 deletions(-) create mode 100644 ggml/src/ggml-feats.h diff --git a/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp index a64ad7a3c74c..13942560895e 100644 --- a/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp +++ b/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp @@ -1,90 +1,11 @@ #include "ggml-backend-impl.h" +#include "ggml-feats.h" -#if defined(__aarch64__) - -#if defined(__linux__) -#include -#elif defined(__APPLE__) -#include -#endif - -#if !defined(HWCAP_FPHP) -#define HWCAP_FPHP (1 << 9) -#endif - -#if !defined(HWCAP_ASIMDHP) -#define HWCAP_ASIMDHP (1 << 10) -#endif - -#if !defined(HWCAP_ASIMDDP) -#define HWCAP_ASIMDDP (1 << 20) -#endif - -#if !defined(HWCAP_SVE) -#define HWCAP_SVE (1 << 22) -#endif - -#if !defined(HWCAP2_SVE2) -#define HWCAP2_SVE2 (1 << 1) -#endif - -#if !defined(HWCAP2_I8MM) -#define HWCAP2_I8MM (1 << 13) -#endif - -#if !defined(HWCAP2_SME) -#define HWCAP2_SME (1 << 23) -#endif - -struct aarch64_features { - // has_neon not needed, aarch64 has NEON guaranteed - bool has_dotprod = false; - bool has_fp16 = false; - bool has_sve = false; - bool has_sve2 = false; - bool has_i8mm = false; - bool has_sme = false; - bool has_sme2 = false; - - aarch64_features() { -#if defined(__linux__) - uint32_t hwcap = getauxval(AT_HWCAP); - uint32_t hwcap2 = getauxval(AT_HWCAP2); - - has_dotprod = !!(hwcap & HWCAP_ASIMDDP); - has_fp16 = !!(hwcap & HWCAP_FPHP) && !!(hwcap & HWCAP_ASIMDHP); - has_sve = !!(hwcap & HWCAP_SVE); - has_sve2 = !!(hwcap2 & HWCAP2_SVE2); - has_i8mm = !!(hwcap2 & HWCAP2_I8MM); - has_sme = !!(hwcap2 & HWCAP2_SME); -#elif defined(__APPLE__) - int oldp = 0; - size_t size = sizeof(oldp); - - if (sysctlbyname("hw.optional.arm.FEAT_DotProd", &oldp, &size, NULL, 0) == 0) { - has_dotprod = static_cast(oldp); - } - - if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &oldp, &size, NULL, 0) == 0) { - has_i8mm = static_cast(oldp); - } - - if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, NULL, 0) == 0) { - has_sme = static_cast(oldp); - } - - if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, NULL, 0) == 0) { - has_sme2 = static_cast(oldp); - } - - // Apple apparently does not implement SVE yet -#endif - } -}; +#if defined(__aarch64__) || defined(_M_ARM64) static int ggml_backend_cpu_aarch64_score() { int score = 1; - aarch64_features af; + ggml_feats_arch64_runtime_t af = ggml_get_aarch64_runtime_features(); #ifdef GGML_USE_DOTPROD if (!af.has_dotprod) { return 0; } @@ -116,4 +37,4 @@ static int ggml_backend_cpu_aarch64_score() { GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_aarch64_score) -# endif // defined(__aarch64__) +# endif // defined(__aarch64__) || defined(_M_ARM64) diff --git a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp index 1c5a459f2190..c1cc33db275e 100644 --- a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp +++ b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp @@ -2,10 +2,12 @@ // SPDX-License-Identifier: MIT // #include -#include -#include +#include +#include +#include #include #include +#include #include #include #include @@ -17,25 +19,21 @@ #include #include #include -#include +#include #include #include +#include +#include #if defined(__linux__) #include +#include #include #include #include #include -#ifndef HWCAP2_SME2 -#define HWCAP2_SME2 (1UL << 37) -#endif #elif defined(__APPLE__) -#include #include #include -#elif defined(_WIN32) -#include -#include #endif #include "kleidiai.h" @@ -43,6 +41,7 @@ #include "ggml-cpu.h" #include "ggml-cpu-impl.h" #include "ggml-impl.h" +#include "ggml-feats.h" #include "ggml-backend-impl.h" #include "ggml-threading.h" #include "traits.h" @@ -64,8 +63,8 @@ struct ggml_kleidiai_context { ggml_kleidiai_kernels * kernels_q4; ggml_kleidiai_kernels * kernels_q8; ggml_kleidiai_kernels * kernels_f32; - int sme_thread_cap; // <= 0 means “SME disabled/unknown”; - int thread_hint; // <= 0 means “no hint” + int sme_thread_cap; // <= 0 means "SME disabled/unknown" + int thread_hint; // <= 0 means "no hint" int chunk_multiplier; } static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, nullptr, 0, -1, 4 }; @@ -93,24 +92,117 @@ static const char* cpu_feature_to_string(cpu_feature f) { } } +#if defined(__linux__) && defined(__aarch64__) +static bool parse_cpu_dir_name(const char* name, size_t* cpu) { + if (strncmp(name, "cpu", 3) != 0 || + name[3] < '0' || name[3] > '9') { + return false; + } + + const char* first = name + 3; + const char* last = name + strlen(name); + + size_t value = 0; + const auto [end, ec] = std::from_chars(first, last, value, 10); + + if (ec != std::errc{} || end != last) { + return false; + } + + *cpu = value; + return true; +} + +static std::vector detect_cpu_ids() { + std::vector cpus; + + DIR * dir = opendir("/sys/devices/system/cpu"); + if (dir == nullptr) { + return cpus; + } + + while (dirent * entry = readdir(dir)) { + size_t cpu = 0; + if (parse_cpu_dir_name(entry->d_name, &cpu)) { + cpus.push_back(cpu); + } + } + closedir(dir); + + std::sort(cpus.begin(), cpus.end()); + cpus.erase(std::unique(cpus.begin(), cpus.end()), cpus.end()); + return cpus; +} +#endif + +#if defined(__APPLE__) && defined(__aarch64__) +static bool apple_sme_counted_perf_level(std::string name) { + for (std::string::size_type i = 0; i < name.size(); ++i) { + name[i] = (char) std::tolower((unsigned char) name[i]); + } + + // Conservative ceiling: only count perf-level names observed to provide full SME throughput. + // Future names should be calibrated here before they raise the automatic SME thread cap. + return name.find("super") != std::string::npos || + name.find("performance") != std::string::npos; +} +#endif + +static void add_smcus_from_smidr(uint64_t smidr, size_t & num_private, std::map & shared_counts) { + // Arm ARM: SMIDR_EL1. SH==0 is implementation-defined; keep the existing + // conservative policy and only treat zero affinity as private. + const uint32_t sh = (uint32_t)((smidr >> 13) & 0x3); + const uint32_t nsmc = (uint32_t)((smidr >> 56) & 0xF); + const size_t shared_count = nsmc == 0xF ? 1 : (size_t)nsmc + 1; + const uint32_t affinity = (uint32_t)(smidr & 0xFFFu); + const uint32_t affinity2 = (uint32_t)((smidr >> 32) & 0xFFFFFu); + const uint32_t id = (affinity2 << 12) | affinity; + + if (nsmc == 0xF) { + GGML_LOG_WARN("kleidiai: NSMC detected as 0xF indicating reseved value, setting min safe shared SMCU count to 1"); + } + + switch (sh) { + case 2: // private SMCU + ++num_private; + break; + case 3: // shared SMCU + if (shared_counts[id] < shared_count) { + shared_counts[id] = shared_count; + } + break; + case 0: + if (id == 0) { + ++num_private; + } else if (shared_counts[id] < shared_count) { + shared_counts[id] = shared_count; + } + break; + default: + break; + } +} + static size_t detect_num_smcus() { - if (!ggml_cpu_has_sme()) { + auto runtime_feat = ggml_get_aarch64_runtime_features(); + if (!runtime_feat.has_sme) { return 0; } #if defined(__linux__) && defined(__aarch64__) // Linux/aarch64: Best-effort count of Streaming Mode Compute Units (SMCUs) via SMIDR_EL1 sysfs. size_t num_private = 0; - std::set shared_ids; + std::map shared_counts; - for (size_t cpu = 0;; ++cpu) { + const std::vector cpus = detect_cpu_ids(); + for (const size_t cpu : cpus) { const std::string path = "/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/regs/identification/smidr_el1"; std::ifstream file(path); if (!file.is_open()) { - break; + continue; } uint64_t smidr = 0; @@ -118,54 +210,69 @@ static size_t detect_num_smcus() { continue; } - // Arm ARM: SMIDR_EL1 - const uint32_t sh = (uint32_t)((smidr >> 13) & 0x3); - // Build an "affinity-like" identifier for shared SMCUs. - // Keep the original packing logic, but isolate it here. - const uint32_t id = (uint32_t)((smidr & 0xFFFu) | ((smidr >> 20) & 0xFFFFF000u)); - - switch (sh) { - case 0b10: // private SMCU - ++num_private; - break; - case 0b11: // shared SMCU - shared_ids.emplace(id); - break; - case 0b00: - // Ambiguous / implementation-defined. Be conservative: - // treat id==0 as private, otherwise as shared. - if (id == 0) ++num_private; - else shared_ids.emplace(id); - break; - default: - break; - } + add_smcus_from_smidr(smidr, num_private, shared_counts); } - return num_private + shared_ids.size(); + size_t total = num_private; + for (const auto & entry : shared_counts) { + total += entry.second; + } + return total; #elif defined(__APPLE__) && defined(__aarch64__) - // table for known M4 variants. Users can override via GGML_KLEIDIAI_SME=. - char chip_name[256] = {}; - size_t size = sizeof(chip_name); - - if (sysctlbyname("machdep.cpu.brand_string", chip_name, &size, nullptr, 0) == 0) { - const std::string brand(chip_name); - - struct ModelSMCU { const char *match; size_t smcus; }; - static const ModelSMCU table[] = { - { "M4 Ultra", 2 }, - { "M4 Max", 2 }, - { "M4 Pro", 2 }, - { "M4", 1 }, - }; + int perf_levels = 0; + size_t size = sizeof(perf_levels); + if (sysctlbyname("hw.nperflevels", &perf_levels, &size, nullptr, 0) != 0 || + size != sizeof(perf_levels) || perf_levels <= 0) { + return 0; + } - for (const auto &e : table) { - if (brand.find(e.match) != std::string::npos) { - return e.smcus; - } + size_t units = 0; + for (int i = 0; i < perf_levels; ++i) { + char key[64] = {}; + int physical_cpus = 0; + int cpus_per_l2 = 0; + + snprintf(key, sizeof(key), "hw.perflevel%d.physicalcpu", i); + size = sizeof(physical_cpus); + if (sysctlbyname(key, &physical_cpus, &size, nullptr, 0) != 0 || + size != sizeof(physical_cpus) || physical_cpus <= 0) { + continue; + } + + snprintf(key, sizeof(key), "hw.perflevel%d.cpusperl2", i); + size = sizeof(cpus_per_l2); + if (sysctlbyname(key, &cpus_per_l2, &size, nullptr, 0) != 0 || + size != sizeof(cpus_per_l2) || cpus_per_l2 <= 0) { + continue; + } + + snprintf(key, sizeof(key), "hw.perflevel%d.name", i); + size = 0; + if (sysctlbyname(key, nullptr, &size, nullptr, 0) != 0 || size == 0) { + continue; + } + + std::string name(size, '\0'); + if (sysctlbyname(key, &name[0], &size, nullptr, 0) != 0) { + continue; + } + name.resize(size); + while (!name.empty() && name.back() == '\0') { + name.pop_back(); + } + + if (apple_sme_counted_perf_level(name)) { + units += (size_t) ((physical_cpus + cpus_per_l2 - 1) / cpus_per_l2); } } + + return units; + +#elif defined(_WIN32) && (defined(_M_ARM64) || defined(__aarch64__)) + // No verified Windows arm64 SMCU detection path yet. Return unknown and use + // GGML_KLEIDIAI_SME=N as a diagnostics/debug override for SME thread cap + // calibration until a detection mechanism is verified on real hardware. return 0; #else @@ -198,15 +305,18 @@ static void init_kleidiai_context(void) { if (!initialized) { initialized = true; + // Optional diagnostics/debug overrides; production defaults come from runtime detection. const char *env_sme = getenv("GGML_KLEIDIAI_SME"); const char *env_threads = getenv("GGML_TOTAL_THREADS"); const char *env_chunk_mult = getenv("GGML_KLEIDIAI_CHUNK_MULTIPLIER"); + auto runtime_feat = ggml_get_aarch64_runtime_features(); + size_t detected_smcus = 0; - ctx.features = (ggml_cpu_has_dotprod() ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) | - (ggml_cpu_has_matmul_int8() ? CPU_FEATURE_I8MM : CPU_FEATURE_NONE) | - ((ggml_cpu_has_sve() && ggml_cpu_get_sve_cnt() == QK8_0) ? CPU_FEATURE_SVE : CPU_FEATURE_NONE); + ctx.features = (runtime_feat.has_dotprod ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) | + (runtime_feat.has_i8mm ? CPU_FEATURE_I8MM : CPU_FEATURE_NONE) | + (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE); if (env_threads) { bool ok = false; @@ -224,54 +334,54 @@ static void init_kleidiai_context(void) { } } - // SME policy: - // - env unset => auto-detect SMCUs; enable SME only if detected > 0. - // - env=0 => force off. - // - env>0 => force N cores, if the binary was built with SME. int sme_cores = 0; bool sme_env_ok = false; bool sme_env_set = (env_sme != nullptr); + const bool has_supported_sme_family = runtime_feat.has_sme; + bool sme_cap_detected = false; + + if (has_supported_sme_family) { + detected_smcus = detect_num_smcus(); + sme_cap_detected = detected_smcus > 0; + // Some platforms expose SME without exposing a calibrated SMCU count. + // Use one SME thread as the conservative default; add platform SMCU detection to raise it. + sme_cores = sme_cap_detected ? (int)detected_smcus : 1; + + if (!sme_env_set && !sme_cap_detected) { + GGML_LOG_INFO("kleidiai: SME detected; SMCU count unavailable, using conservative SME thread cap=1\n"); + } + } + + // Runtime-detect SME support and available SMCUs first. The detected SMCU + // count is used as the SME thread cap, and GGML_KLEIDIAI_SME can debug-override that: + // - unset: use runtime detection. + // - 0: disable SME-family kernels. + // - N > 0: use N as the SME thread cap, if an SME-family kernel is selectable. if (sme_env_set) { bool ok = false; int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok); sme_env_ok = ok; - if (!ok) { - GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n"); - detected_smcus = detect_num_smcus(); - sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0; - } else if (v == 0) { - sme_cores = 0; - } else if (!ggml_cpu_has_sme()) { - GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but the binary was not built with SME; disabling SME\n", v); - sme_cores = 0; + if (ok) { + if (has_supported_sme_family) { + sme_cores = v; + } else { + if (v > 0) { + GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but SME is not supported on this CPU; disabling SME-family kernels\n", v); + } + sme_cores = 0; + } } else { - sme_cores = v; + GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; using automatic SME thread cap\n"); } - } else { - detected_smcus = detect_num_smcus(); - sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0; } - if (!sme_env_set && ggml_cpu_has_sme() && sme_cores == 0) { - GGML_LOG_WARN("kleidiai: runtime SME-core detection returned 0; falling back to NEON\n"); - } - - if (sme_cores > 0) { + if (sme_cores > 0 && has_supported_sme_family) { ctx.features |= CPU_FEATURE_SME; -#if defined(__aarch64__) && defined(__linux__) - // ARM guarantees SME2 implies SME, so only check SME2 when SME is enabled. - if (getauxval(AT_HWCAP2) & HWCAP2_SME2) { + if (runtime_feat.has_sme2) { ctx.features |= CPU_FEATURE_SME2; } -#elif defined(__aarch64__) && defined(__APPLE__) - int feat_sme2 = 0; - size_t size = sizeof(feat_sme2); - if (sysctlbyname("hw.optional.arm.FEAT_SME2", &feat_sme2, &size, NULL, 0) == 0 && feat_sme2) { - ctx.features |= CPU_FEATURE_SME2; - } -#endif } // Kernel selection @@ -297,16 +407,19 @@ static void init_kleidiai_context(void) { GGML_LOG_INFO("kleidiai: primary f32 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_f32->required_cpu)); } - ctx.sme_thread_cap = (ctx.features & CPU_FEATURE_SME) ? sme_cores : 0; + const bool has_selected_sme_family_kernel = + (ctx.kernels_q4 && is_sme_family(ctx.kernels_q4->required_cpu)) || + (ctx.kernels_q8 && is_sme_family(ctx.kernels_q8->required_cpu)) || + (ctx.kernels_f32 && is_sme_family(ctx.kernels_f32->required_cpu)); + ctx.sme_thread_cap = has_selected_sme_family_kernel ? sme_cores : 0; - if (ctx.features & CPU_FEATURE_SME) { - const bool has_sme2 = (ctx.features & CPU_FEATURE_SME2) != CPU_FEATURE_NONE; + if (has_selected_sme_family_kernel) { if (sme_env_set && sme_env_ok && sme_cores > 0) { - GGML_LOG_INFO("kleidiai: SME%s enabled (GGML_KLEIDIAI_SME=%d override)\n", - has_sme2 ? "2" : "", sme_cores); + GGML_LOG_INFO("kleidiai: SME enabled (GGML_KLEIDIAI_SME=%d debug override)\n", sme_cores); + } else if (sme_cap_detected) { + GGML_LOG_INFO("kleidiai: SME enabled (runtime-detected SME thread cap=%d)\n", sme_cores); } else { - GGML_LOG_INFO("kleidiai: SME%s enabled (runtime-detected SME cores=%d)\n", - has_sme2 ? "2" : "", sme_cores); + GGML_LOG_INFO("kleidiai: SME enabled (runtime SME detected, conservative thread cap=%d)\n", sme_cores); } } else { GGML_LOG_INFO("kleidiai: SME disabled\n"); @@ -467,7 +580,7 @@ static int kleidiai_collect_kernel_chain_common( } if (is_sme_family(primary->required_cpu)) { - const cpu_feature fallback_mask = static_cast(features & ~CPU_FEATURE_SME & ~CPU_FEATURE_SME2); + const cpu_feature fallback_mask = static_cast(features & ~(CPU_FEATURE_SME | CPU_FEATURE_SME2)); if (fallback_mask != CPU_FEATURE_NONE) { ggml_kleidiai_kernels * fallback = select_fallback(fallback_mask); if (fallback && fallback != primary && @@ -1077,13 +1190,14 @@ class tensor_traits : public ggml::cpu::tensor_traits { const int ith_total = params->ith; int sme_slot = -1; + int non_sme_slot = -1; for (int i = 0; i < runtime_count; ++i) { if (is_sme_family(runtime[i].kernels->required_cpu)) { sme_slot = i; break; } } - int non_sme_slot = -1; + for (int i = 0; i < runtime_count; ++i) { if (!is_sme_family(runtime[i].kernels->required_cpu)) { non_sme_slot = i; diff --git a/ggml/src/ggml-feats.h b/ggml/src/ggml-feats.h new file mode 100644 index 000000000000..d4a8c83a7fd8 --- /dev/null +++ b/ggml/src/ggml-feats.h @@ -0,0 +1,166 @@ +#pragma once + +#if defined(__aarch64__) || defined(_M_ARM64) + +#if defined(__linux__) +#include +#include + +#if !defined(HWCAP2_SVE2) +#define HWCAP2_SVE2 (1ULL << 1) +#endif + +#if !defined(HWCAP_FPHP) +#define HWCAP_FPHP (1 << 9) +#endif + +#if !defined(HWCAP_ASIMDHP) +#define HWCAP_ASIMDHP (1 << 10) +#endif + +#if !defined(HWCAP2_I8MM) +#define HWCAP2_I8MM (1ULL << 13) +#endif + +#if !defined(HWCAP_ASIMDDP) +#define HWCAP_ASIMDDP (1 << 20) +#endif + +#if !defined(HWCAP_SVE) +#define HWCAP_SVE (1 << 22) +#endif + +#if !defined(HWCAP2_SME) +#define HWCAP2_SME (1ULL << 23) +#endif + +#if !defined(HWCAP2_SME2) +#define HWCAP2_SME2 (1ULL << 37) +#endif + +#if !defined(PR_SVE_GET_VL) +#define PR_SVE_GET_VL 51 +#endif + +#if !defined(PR_SVE_VL_LEN_MASK) +#define PR_SVE_VL_LEN_MASK 0xffff +#endif + +#elif defined(__APPLE__) +#include +#elif defined(_WIN32) +#include + +#if !defined(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE 43 +#endif + +#if !defined(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_SVE_INSTRUCTIONS_AVAILABLE 46 +#endif + +#if !defined(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47 +#endif + +#if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66 +#endif + +#if !defined(PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE 67 +#endif + +#if !defined(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_SME_INSTRUCTIONS_AVAILABLE 70 +#endif + +#if !defined(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) +#define PF_ARM_SME2_INSTRUCTIONS_AVAILABLE 71 +#endif + +#endif + +typedef struct ggml_feats_arch64_runtime { + bool has_dotprod; + bool has_fp16; + bool has_sve; + bool has_sve2; + bool has_i8mm; + bool has_sme; + bool has_sme2; + int sve_cnt; +} ggml_feats_arch64_runtime_t; + +static inline ggml_feats_arch64_runtime_t ggml_get_aarch64_runtime_features(void) { + ggml_feats_arch64_runtime_t runtime_feat = {}; + +#if defined(__linux__) + const unsigned long hwcap = getauxval(AT_HWCAP); + const unsigned long hwcap2 = getauxval(AT_HWCAP2); + + runtime_feat.has_dotprod = !!(hwcap & HWCAP_ASIMDDP); + runtime_feat.has_fp16 = !!(hwcap & HWCAP_FPHP) && !!(hwcap & HWCAP_ASIMDHP);; + runtime_feat.has_sve = !!(hwcap & HWCAP_SVE); + runtime_feat.has_sve2 = !!(hwcap2 & HWCAP2_SVE2); + runtime_feat.has_i8mm = !!(hwcap2 & HWCAP2_I8MM); + runtime_feat.has_sme = !!(hwcap2 & HWCAP2_SME); + runtime_feat.has_sme2 = !!(hwcap2 & HWCAP2_SME2); + + if (runtime_feat.has_sve) { + const int vl = prctl(PR_SVE_GET_VL); + if (vl >= 0) { + runtime_feat.sve_cnt = vl & PR_SVE_VL_LEN_MASK; + } + } +#elif defined(__APPLE__) + int oldp = 0; + size_t size = sizeof(oldp); + + if (sysctlbyname("hw.optional.arm.FEAT_DotProd", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_dotprod = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_FP16", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_fp16 = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_SVE", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_sve = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_SVE2", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_sve2 = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_i8mm = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_sme = static_cast(oldp); + } + + if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, nullptr, 0) == 0) { + runtime_feat.has_sme2 = static_cast(oldp); + } + + // Apple does not support userspace non-streaming SVE; keep SVE vector length unknown. + runtime_feat.sve_cnt = 0; +#elif defined (_WIN32) + runtime_feat.has_dotprod = IsProcessorFeaturePresent(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_fp16 = IsProcessorFeaturePresent(PF_ARM_V82_FP16_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_sve = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_sve2 = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_i8mm = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_sme = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0; + runtime_feat.has_sme2 = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0; + + // Windows exposes SVE feature presence, but not the runtime SVE vector length here. + runtime_feat.sve_cnt = 0; +#endif + + return runtime_feat; +} + +#endif // defined(__aarch64__) || defined(_M_ARM64) From d8a8beac22d450ebadf175a8ce7b6bf49b66db14 Mon Sep 17 00:00:00 2001 From: HarrisonSec Date: Wed, 12 Aug 2026 05:07:48 -0700 Subject: [PATCH 038/211] gguf : harden loader against malformed tensor dims and metadata types (#25596) * gguf : harden loader against malformed tensor dims and metadata types * gguf: address review on malformed-metadata hardening - report the expected vs. actual type when general.alignment is not u32 - use ggml_nelements() > 0 for the zero-element guard and keep the representability checks visually aligned - add test-gguf cases for a wrong-typed alignment key and a zero-dim tensor (both used to crash: assert-abort and SIGFPE respectively) Ran tests/test-gguf: 164/164 pass. Used an AI assistant to help draft these edits; reviewed and verified by me. * cont : less comments Co-authored-by: Georgi Gerganov --------- Co-authored-by: Georgi Gerganov --- ggml/src/gguf.cpp | 15 ++++++++++++--- tests/test-gguf.cpp | 38 ++++++++++++++++++++++++++++++-------- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp index 9f9e4fe5d104..6c7b5817812b 100644 --- a/ggml/src/gguf.cpp +++ b/ggml/src/gguf.cpp @@ -611,6 +611,13 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr GGML_ASSERT(int64_t(ctx->kv.size()) == n_kv); const int alignment_idx = gguf_find_key(ctx, GGUF_KEY_GENERAL_ALIGNMENT); + if (alignment_idx != -1 && gguf_get_kv_type(ctx, alignment_idx) != GGUF_TYPE_UINT32) { + GGML_LOG_ERROR("%s: key '%s' must be of type %s but is %s\n", + __func__, GGUF_KEY_GENERAL_ALIGNMENT, gguf_type_name(GGUF_TYPE_UINT32), + gguf_type_name(gguf_get_kv_type(ctx, alignment_idx))); + gguf_free(ctx); + return nullptr; + } ctx->alignment = alignment_idx == -1 ? GGUF_DEFAULT_ALIGNMENT : gguf_get_val_u32(ctx, alignment_idx); if (ctx->alignment == 0 || (ctx->alignment & (ctx->alignment - 1)) != 0) { @@ -682,9 +689,11 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr } // check that the total number of elements is representable - if (ok && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || - (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) || - (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) { + // (a zero-element tensor is trivially representable; the guard also avoids a division by zero below) + if (ok && ggml_nelements(&info.t) > 0 && + ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || + (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) || + (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) { GGML_LOG_ERROR("%s: total number of elements in tensor '%s' with shape " "(%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") is >= %" PRIi64 "\n", diff --git a/tests/test-gguf.cpp b/tests/test-gguf.cpp index 2875dec806da..fc636186f4c5 100644 --- a/tests/test-gguf.cpp +++ b/tests/test-gguf.cpp @@ -31,11 +31,13 @@ enum handcrafted_file_type { // HANDCRAFTED_KV_BAD_VALUE_SIZE = 30 + offset_has_kv, // removed because it can result in allocations > 1 TB (default sanitizer limit) HANDCRAFTED_KV_DUPLICATE_KEY = 40 + offset_has_kv, HANDCRAFTED_KV_BAD_ALIGN = 50 + offset_has_kv, + HANDCRAFTED_KV_WRONG_TYPE_ALIGN = 55 + offset_has_kv, HANDCRAFTED_KV_SUCCESS = 800 + offset_has_kv, HANDCRAFTED_TENSORS_BAD_NAME_SIZE = 10 + offset_has_tensors, HANDCRAFTED_TENSORS_BAD_N_DIMS = 20 + offset_has_tensors, HANDCRAFTED_TENSORS_BAD_SHAPE = 30 + offset_has_tensors, + HANDCRAFTED_TENSORS_ZERO_DIM = 35 + offset_has_tensors, HANDCRAFTED_TENSORS_NE_TOO_BIG = 40 + offset_has_tensors, HANDCRAFTED_TENSORS_NBYTES_TOO_BIG = 45 + offset_has_tensors, HANDCRAFTED_TENSORS_BAD_TYPE = 50 + offset_has_tensors, @@ -69,11 +71,13 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h case HANDCRAFTED_KV_BAD_TYPE: return "KV_BAD_TYPE"; case HANDCRAFTED_KV_DUPLICATE_KEY: return "KV_DUPLICATE_KEY"; case HANDCRAFTED_KV_BAD_ALIGN: return "KV_BAD_ALIGN"; + case HANDCRAFTED_KV_WRONG_TYPE_ALIGN: return "KV_WRONG_TYPE_ALIGN"; case HANDCRAFTED_KV_SUCCESS: return "KV_RANDOM_KV"; case HANDCRAFTED_TENSORS_BAD_NAME_SIZE: return "TENSORS_BAD_NAME_SIZE"; case HANDCRAFTED_TENSORS_BAD_N_DIMS: return "TENSORS_BAD_N_DIMS"; case HANDCRAFTED_TENSORS_BAD_SHAPE: return "TENSORS_BAD_SHAPE"; + case HANDCRAFTED_TENSORS_ZERO_DIM: return "TENSORS_ZERO_DIM"; case HANDCRAFTED_TENSORS_NE_TOO_BIG: return "TENSORS_NE_TOO_BIG"; case HANDCRAFTED_TENSORS_NBYTES_TOO_BIG: return "TENSORS_NBYTES_TOO_BIG"; case HANDCRAFTED_TENSORS_BAD_TYPE: return "TENSORS_BAD_TYPE"; @@ -95,6 +99,9 @@ static std::string handcrafted_file_type_name(const enum handcrafted_file_type h } static bool expect_context_not_null(const enum handcrafted_file_type hft) { + if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) { + return true; + } if (hft < offset_has_kv) { return hft >= HANDCRAFTED_HEADER_EMPTY; } @@ -257,9 +264,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft } { uint64_t n_kv = kv_types.size(); - if (hft == HANDCRAFTED_KV_BAD_ALIGN || - hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN || - hft == HANDCRAFTED_DATA_BAD_ALIGN || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) { + if (hft == HANDCRAFTED_KV_BAD_ALIGN || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN || + hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN || + hft == HANDCRAFTED_DATA_BAD_ALIGN || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) { n_kv += 1; } else if (hft == HANDCRAFTED_HEADER_BAD_N_KV) { @@ -344,15 +351,17 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft helper_write(file, data, hft == HANDCRAFTED_KV_BAD_TYPE ? 1 : gguf_type_size(type)); } - if (hft == HANDCRAFTED_KV_BAD_ALIGN || - hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN || - hft == HANDCRAFTED_DATA_BAD_ALIGN || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) { + if (hft == HANDCRAFTED_KV_BAD_ALIGN || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN || + hft == HANDCRAFTED_TENSORS_BAD_ALIGN || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN || + hft == HANDCRAFTED_DATA_BAD_ALIGN || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) { const uint64_t n = strlen(GGUF_KEY_GENERAL_ALIGNMENT); helper_write(file, n); helper_write(file, GGUF_KEY_GENERAL_ALIGNMENT, n); - const int32_t type = gguf_type(GGUF_TYPE_UINT32); + // HANDCRAFTED_KV_WRONG_TYPE_ALIGN declares general.alignment with a non-UINT32 type, + // which the loader must reject cleanly instead of aborting on an assertion + const int32_t type = hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ? int32_t(GGUF_TYPE_INT32) : int32_t(GGUF_TYPE_UINT32); helper_write(file, type); alignment = expect_context_not_null(hft) ? 1 : 13; @@ -403,6 +412,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft break; } } + if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) { + n_dims = 2; + } if (hft == HANDCRAFTED_TENSORS_BAD_N_DIMS) { const uint32_t n_dims_bad = GGML_MAX_DIMS + 1; helper_write(file, n_dims_bad); @@ -415,6 +427,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft for (uint32_t j = 0; j < n_dims; ++j) { helper_write(file, bad_dim); } + } else if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) { + const int64_t zero_shape[2] = { shape[0], 0 }; + helper_write(file, zero_shape, 2*sizeof(int64_t)); } else if (hft == HANDCRAFTED_TENSORS_NE_TOO_BIG){ const int64_t big_dim = 4*int64_t(INT32_MAX); for (uint32_t j = 0; j < n_dims; ++j) { @@ -446,6 +461,9 @@ static FILE * get_handcrafted_file(const unsigned int seed, const enum handcraft for (uint32_t i = 1; i < n_dims; ++i) { ne *= shape[i]; } + if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) { + ne = 0; + } offset += GGML_PAD(ggml_row_size(type, ne), (uint64_t) alignment); } @@ -747,11 +765,13 @@ static std::pair test_handcrafted_file(const unsigned int seed) { HANDCRAFTED_KV_BAD_TYPE, HANDCRAFTED_KV_DUPLICATE_KEY, HANDCRAFTED_KV_BAD_ALIGN, + HANDCRAFTED_KV_WRONG_TYPE_ALIGN, HANDCRAFTED_KV_SUCCESS, HANDCRAFTED_TENSORS_BAD_NAME_SIZE, HANDCRAFTED_TENSORS_BAD_N_DIMS, HANDCRAFTED_TENSORS_BAD_SHAPE, + HANDCRAFTED_TENSORS_ZERO_DIM, HANDCRAFTED_TENSORS_NE_TOO_BIG, HANDCRAFTED_TENSORS_NBYTES_TOO_BIG, HANDCRAFTED_TENSORS_BAD_TYPE, @@ -840,7 +860,9 @@ static std::pair test_handcrafted_file(const unsigned int seed) { ntest++; } - if (expect_context_not_null(hft) && hft >= offset_has_tensors) { + // HANDCRAFTED_TENSORS_ZERO_DIM deliberately mangles the tensor shapes to 0 elements, + // so only assert that it loads without crashing; skip the exact-geometry comparison. + if (expect_context_not_null(hft) && hft >= offset_has_tensors && hft != HANDCRAFTED_TENSORS_ZERO_DIM) { printf("%s: - check_tensors: ", __func__); if (handcrafted_check_tensors(gguf_ctx, seed)) { printf("\033[1;32mOK\033[0m\n"); From 680a9ae63d60d35c21a0dcd7d3fabdb9c6bfc963 Mon Sep 17 00:00:00 2001 From: Daniel Bevenius Date: Wed, 12 Aug 2026 14:15:03 +0200 Subject: [PATCH 039/211] cmake : introduce semantic versioning (#26839) * cmake : introduce semantic versioning (wip) This commit introduces semantic versioning to llama.cpp. * squash! cmake : introduce semantic versioning (wip) * cmake : update test-cmake README notes [no ci] * include libmtmd in output so show its semversioned * ci : add make-release workflow * ci : fix build number check in build-cmake-pkg.yml * examples : remove trailing whitespace * ci : abort if upstream ggml version does not exist * ci : extract step contents into scripts * ci : add GGML_NATIVE=OFF to ubuntu job * examples : remove CI build information from test-cmake [no ci] This commit removes the nightly/release information that I added previously to keep this focused only on using building and installing llama.cpp with cmake and being able to quickly verify changes or troubleshoot issues. * ci : merge scripts into single script * remove -dev-build_number support This commit removes the incremental build number (versioning) support that I added. This was incorrect and we should only use the semver for the version. Releases will be tag a nightly build and package maintainers/managers that build from source can use the tag and it is therefor important that the correct version is reported. So a nightly-build will report the semver without the build number. The build number and commit as availble via cmake and test-cmake has been updated to include an example of using them: ```console $ ./build.sh [test-cmake] version: 0.1.0, build: 10360 (08c69e381) ... ``` Refs: https://github.com/ggml-org/llama.cpp/pull/26839#discussion_r3755836969 * docs: add initial release.md documentation * cmake : clean-up and add LLAMA_BUILD_IS_DEV option * ci : remove version input from make-release job * ci : add LLAMA_BUILD_IS_DEV=OFF to build-cmake-pkg.yml Refs: https://github.com/danbev/llama.cpp/actions/runs/31576801921/job/94050639145 * docs : update release notes with LLAMA_BUILD_IS_DEV info [no ci] * ci : add TODO to winget workflow [no ci] --------- Co-authored-by: Georgi Gerganov --- .github/workflows/build-cmake-pkg.yml | 8 +- .github/workflows/build-cpu.yml | 3 +- .github/workflows/make-release.yml | 46 ++++++++++ .github/workflows/winget.yml | 2 + CMakeLists.txt | 30 +++++-- README.md | 1 + app/llama.cpp | 8 +- cmake/llama-config.cmake.in | 2 +- cmake/llama.pc.in | 2 +- common/CMakeLists.txt | 4 +- common/arg.cpp | 3 +- common/build-info.cpp.in | 6 +- common/build-info.h | 2 +- docs/release.md | 49 +++++++++++ examples/test-cmake/.gitignore | 3 + examples/test-cmake/CMakeLists.txt | 13 +++ examples/test-cmake/README.md | 36 ++++++++ examples/test-cmake/build-install.sh | 19 +++++ examples/test-cmake/build.sh | 7 ++ examples/test-cmake/test-cmake.cpp | 12 +++ include/llama.h | 2 + scripts/make-release-checks.sh | 83 +++++++++++++++++++ src/CMakeLists.txt | 9 +- src/llama.cpp | 4 + tests/test-quantize-stats.cpp | 2 +- tools/cvector-generator/cvector-generator.cpp | 2 +- tools/gguf-split/gguf-split.cpp | 2 +- tools/mtmd/CMakeLists.txt | 4 +- tools/quantize/quantize.cpp | 2 +- 29 files changed, 336 insertions(+), 30 deletions(-) create mode 100644 .github/workflows/make-release.yml create mode 100644 docs/release.md create mode 100644 examples/test-cmake/.gitignore create mode 100644 examples/test-cmake/CMakeLists.txt create mode 100644 examples/test-cmake/README.md create mode 100755 examples/test-cmake/build-install.sh create mode 100755 examples/test-cmake/build.sh create mode 100644 examples/test-cmake/test-cmake.cpp create mode 100755 scripts/make-release-checks.sh diff --git a/.github/workflows/build-cmake-pkg.yml b/.github/workflows/build-cmake-pkg.yml index 5becff09c1bc..a958870de3d6 100644 --- a/.github/workflows/build-cmake-pkg.yml +++ b/.github/workflows/build-cmake-pkg.yml @@ -21,6 +21,7 @@ jobs: -DLLAMA_BUILD_TOOLS=OFF \ -DLLAMA_BUILD_EXAMPLES=OFF \ -DLLAMA_BUILD_APP=OFF \ + -DLLAMA_BUILD_IS_DEV=OFF \ -DCMAKE_BUILD_TYPE=Release cmake --build build --config Release cmake --install build --prefix "$PREFIX" --config Release @@ -29,7 +30,12 @@ jobs: tclsh <<'EOF' set build(commit) [string trim [exec git rev-parse --short HEAD]] set build(number) [string trim [exec git rev-list --count HEAD]] - set build(version) "0.0.$build(number)" + + set cmakelists [read [open "CMakeLists.txt" r]] + regexp {set\(LLAMA_VERSION_MAJOR\s+(\d+)\)} $cmakelists -> major + regexp {set\(LLAMA_VERSION_MINOR\s+(\d+)\)} $cmakelists -> minor + regexp {set\(LLAMA_VERSION_PATCH\s+(\d+)\)} $cmakelists -> patch + set build(version) "$major.$minor.$patch" set llamaconfig [read [open "$env(LLAMA_CONFIG)" r]] set checks [list "set\\(LLAMA_VERSION \\s+$build(version)\\)" \ diff --git a/.github/workflows/build-cpu.yml b/.github/workflows/build-cpu.yml index 30b07ce7882c..57622e1d4507 100644 --- a/.github/workflows/build-cpu.yml +++ b/.github/workflows/build-cpu.yml @@ -95,7 +95,8 @@ jobs: run: | cmake -B build \ -DLLAMA_FATAL_WARNINGS=ON \ - -DGGML_RPC=ON + -DGGML_RPC=ON \ + -DGGML_NATIVE=OFF time cmake --build build --config Release -j $(nproc) - name: Test diff --git a/.github/workflows/make-release.yml b/.github/workflows/make-release.yml new file mode 100644 index 000000000000..fed9c877c71e --- /dev/null +++ b/.github/workflows/make-release.yml @@ -0,0 +1,46 @@ +name: Make Release + +on: + workflow_dispatch: + inputs: + dry_run: + description: 'Dry run - validate without creating the tag' + required: true + type: boolean + default: true + +env: + GH_TOKEN: ${{ github.token }} + +permissions: + contents: write + +jobs: + make-release: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: Run release checks + id: checks + run: bash scripts/make-release-checks.sh ${{ github.event.inputs.dry_run == 'true' && '--dry-run' || '' }} + env: + GITHUB_REPOSITORY: ${{ github.repository }} + + - name: Create release tag + if: ${{ github.event.inputs.dry_run == 'false' }} + run: | + VERSION="${{ steps.checks.outputs.version }}" + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git tag -a "${VERSION}" -m "Release ${VERSION}" + git push origin "${VERSION}" + echo "Created and pushed tag ${VERSION}" + + - name: Dry run summary + if: ${{ github.event.inputs.dry_run == 'true' }} + run: | + echo "Dry run complete - all checks passed." + echo "Would have created tag: ${{ steps.checks.outputs.version }}" diff --git a/.github/workflows/winget.yml b/.github/workflows/winget.yml index 69e24f940099..c0a814f3adbf 100644 --- a/.github/workflows/winget.yml +++ b/.github/workflows/winget.yml @@ -19,6 +19,8 @@ jobs: run: | cargo binstall komac@2.16.0 -y + # TODO: This should later be updated to publish releases instead of + # development release builds. - name: Find latest release id: find_latest_release uses: actions/github-script@v8 diff --git a/CMakeLists.txt b/CMakeLists.txt index 3df1d82dbe09..b2092d12ddd4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,6 +2,26 @@ cmake_minimum_required(VERSION 3.14...3.28) # for add_link_options and implicit project("llama.cpp" C CXX) include(CheckIncludeFileCXX) +### llama.cpp version +set(LLAMA_VERSION_MAJOR 0) +set(LLAMA_VERSION_MINOR 1) +set(LLAMA_VERSION_PATCH 0) +set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}") + +# whether this is a development/nightly build +# set this to OFF when making a release from a release tag (vX.Y.Z) +# ref: https://github.com/ggml-org/ggml/discussions/1579 +option(LLAMA_BUILD_IS_DEV "llama: dev build" ON) + +if (LLAMA_BUILD_IS_DEV) + set(LLAMA_VERSION "${LLAMA_VERSION_BASE}-dev") +else() + # TODO: check that the current commit is tagged correctly according to the version specified above + set(LLAMA_VERSION "${LLAMA_VERSION_BASE}") +endif() + +message(STATUS "llama.cpp version: ${LLAMA_VERSION}") + #set(CMAKE_WARN_DEPRECATED YES) set(CMAKE_WARN_UNUSED_CLI YES) @@ -24,9 +44,6 @@ if (CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) set(LLAMA_STANDALONE ON) include(git-vars) - - # configure project version - # TODO else() set(LLAMA_STANDALONE OFF) endif() @@ -139,7 +156,6 @@ endif() if (NOT DEFINED LLAMA_BUILD_COMMIT) set(LLAMA_BUILD_COMMIT ${BUILD_COMMIT}) endif() -set(LLAMA_INSTALL_VERSION 0.0.${LLAMA_BUILD_NUMBER}) # override ggml options set(GGML_ALL_WARNINGS ${LLAMA_ALL_WARNINGS}) @@ -275,12 +291,12 @@ configure_package_config_file( LLAMA_BIN_INSTALL_DIR ) write_basic_package_version_file( - ${CMAKE_CURRENT_BINARY_DIR}/llama-version.cmake - VERSION ${LLAMA_INSTALL_VERSION} + ${CMAKE_CURRENT_BINARY_DIR}/llama-config-version.cmake + VERSION ${LLAMA_VERSION} COMPATIBILITY SameMajorVersion) install(FILES ${CMAKE_CURRENT_BINARY_DIR}/llama-config.cmake - ${CMAKE_CURRENT_BINARY_DIR}/llama-version.cmake + ${CMAKE_CURRENT_BINARY_DIR}/llama-config-version.cmake DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/llama) configure_file(cmake/llama.pc.in diff --git a/README.md b/README.md index fc7633a7c550..1b341e7fbc6e 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,7 @@ The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-or - [XCFramework](docs/xcframework.md) - [Completions](docs/completions.md) - [Models](docs/models.md) +- [Release process](docs/release.md) ## Contributing diff --git a/app/llama.cpp b/app/llama.cpp index 2cf1aa876ce0..3b7e46f20dee 100644 --- a/app/llama.cpp +++ b/app/llama.cpp @@ -1,5 +1,7 @@ #include "build-info.h" +#include "llama.h" + #include #include #include @@ -77,12 +79,12 @@ static const command cmds[] = { #undef UPDATE_HIDDEN -static int version(int argc, char ** argv) { - printf("%s\n", llama_build_info()); +static int version(int /*argc*/, char ** /*argv*/) { + llama_print_build_info(llama_version()); return 0; } -static int licenses(int argc, char ** argv) { +static int licenses(int /*argc*/, char ** /*argv*/) { for (int i = 0; LICENSES[i]; ++i) { printf("%s\n", LICENSES[i]); } diff --git a/cmake/llama-config.cmake.in b/cmake/llama-config.cmake.in index b4defc76ff0e..6db73577ae6d 100644 --- a/cmake/llama-config.cmake.in +++ b/cmake/llama-config.cmake.in @@ -1,4 +1,4 @@ -set(LLAMA_VERSION @LLAMA_INSTALL_VERSION@) +set(LLAMA_VERSION @LLAMA_VERSION@) set(LLAMA_BUILD_COMMIT @LLAMA_BUILD_COMMIT@) set(LLAMA_BUILD_NUMBER @LLAMA_BUILD_NUMBER@) set(LLAMA_SHARED_LIB @BUILD_SHARED_LIBS@) diff --git a/cmake/llama.pc.in b/cmake/llama.pc.in index 6fb58b5f6881..31b043c0e399 100644 --- a/cmake/llama.pc.in +++ b/cmake/llama.pc.in @@ -5,6 +5,6 @@ includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@ Name: llama Description: Port of Facebook's LLaMA model in C/C++ -Version: @LLAMA_INSTALL_VERSION@ +Version: @LLAMA_VERSION@ Libs: -L${libdir} -lggml -lggml-base -lllama Cflags: -I${includedir} diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 799d227519f9..d6cfc9a00872 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -121,8 +121,8 @@ add_library(${TARGET} ) set_target_properties(${TARGET} PROPERTIES - VERSION ${LLAMA_INSTALL_VERSION} - SOVERSION 0 + VERSION ${LLAMA_VERSION_BASE} + SOVERSION ${LLAMA_VERSION_MAJOR} MACHO_CURRENT_VERSION 0 # keep macOS linker from seeing oversized version number ) diff --git a/common/arg.cpp b/common/arg.cpp index cb314eee7044..b2fddfe5f2f0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1390,8 +1390,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--version"}, "show version and build info", [](common_params &) { - fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit()); - fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target()); + llama_print_build_info(llama_version()); exit(0); } )); diff --git a/common/build-info.cpp.in b/common/build-info.cpp.in index f888fd079fa5..4ec3397081b4 100644 --- a/common/build-info.cpp.in +++ b/common/build-info.cpp.in @@ -29,7 +29,7 @@ const char * llama_build_info(void) { return s.c_str(); } -void llama_print_build_info(void) { - fprintf(stderr, "%s: build = %d (%s)\n", __func__, llama_build_number(), llama_commit()); - fprintf(stderr, "%s: built with %s for %s\n", __func__, llama_compiler(), llama_build_target()); +void llama_print_build_info(const char * llama_version) { + fprintf(stderr, "version: %s (build %d, commit %s)\n", llama_version, llama_build_number(), llama_commit()); + fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target()); } diff --git a/common/build-info.h b/common/build-info.h index 382cfa78500a..1e564591a612 100644 --- a/common/build-info.h +++ b/common/build-info.h @@ -8,4 +8,4 @@ const char * llama_compiler(void); const char * llama_build_target(void); const char * llama_build_info(void); -void llama_print_build_info(void); +void llama_print_build_info(const char *); diff --git a/docs/release.md b/docs/release.md new file mode 100644 index 000000000000..4335ef9d429b --- /dev/null +++ b/docs/release.md @@ -0,0 +1,49 @@ +# Release process + +llama.cpp uses [semantic versioning](https://semver.org) (`MAJOR.MINOR.PATCH`). + +## Version bump guidelines + +| Change type | Version component | +|---|---| +| Breaking change to the public C API (`include/llama.h`) | `MAJOR` | +| Backward-compatible features, model support, or API addition | `MINOR` | +| Bug fix with no API change | `PATCH` | + +The version is set in the three variables at the top of the root `CMakeLists.txt`: + +```cmake +set(LLAMA_VERSION_MAJOR 0) +set(LLAMA_VERSION_MINOR 1) +set(LLAMA_VERSION_PATCH 0) +``` + +_A version bump should be included in the PR that introduces the change, or in a +dedicated bump commit merged before the release is cut._ + +_TODO: add PR labels (`semver: patch`, `semver: minor`, `semver: major`) to help +identify which PRs require a version bump before cutting a release._ + +## Making a release + +Releases are created by running the [make-release](.github/workflows/make-release.yml) +which is a manual workflow. + +The workflow creates an annotated git tag (e.g. `v0.1.0`) and pushes it to the +remote. No GitHub Release object is created, the tag is the release artifact. + +## Building a release + +By default, `LLAMA_BUILD_IS_DEV=ON` which appends a `-dev` suffix to `LLAMA_VERSION`, +marking the build as a nightly/development build. Distributors building from a +release tag must pass `-DLLAMA_BUILD_IS_DEV=OFF` to produce a clean version string +(e.g. `0.1.0` instead of `0.1.0-dev`). + +## How releases reach users +Currently releases are not published to github releases, only nightly/development +builds are available there. The way users can access releases are using the following +channels: + +- **llama-install.sh** — downloads pre-built binaries built from the release tag. +- **Package managers** — consume the git tag directly. +- **Build from source** — users clone the repo and check out the tag. diff --git a/examples/test-cmake/.gitignore b/examples/test-cmake/.gitignore new file mode 100644 index 000000000000..0ddff317a4b5 --- /dev/null +++ b/examples/test-cmake/.gitignore @@ -0,0 +1,3 @@ +llama-build-install +install +build diff --git a/examples/test-cmake/CMakeLists.txt b/examples/test-cmake/CMakeLists.txt new file mode 100644 index 000000000000..ed5cb1f3c262 --- /dev/null +++ b/examples/test-cmake/CMakeLists.txt @@ -0,0 +1,13 @@ +cmake_minimum_required(VERSION 3.14) +project(llama-simple) + +set(CMAKE_CXX_STANDARD 17) + +find_package(llama 0.1.0 REQUIRED) + +add_executable(test-cmake test-cmake.cpp) +target_link_libraries(test-cmake PRIVATE llama) +target_compile_definitions(test-cmake PRIVATE + LLAMA_BUILD_NUMBER=${LLAMA_BUILD_NUMBER} + LLAMA_BUILD_COMMIT="${LLAMA_BUILD_COMMIT}" +) diff --git a/examples/test-cmake/README.md b/examples/test-cmake/README.md new file mode 100644 index 000000000000..21e5eb9607dd --- /dev/null +++ b/examples/test-cmake/README.md @@ -0,0 +1,36 @@ +## cmake-test + +This is just for manually testing/developing of a llama.cpp installation to +enable troubleshooting issues and exploration. The idea is that this can be used +after making changes to llama.cpp installation cmake configuration and then +verify it locally. + +### Usage +The following will configure, build, and install llama.cpp + +Configuring/build/install: +```console +./build-install.sh +``` +The above command will create a directory named `install` in the current directory +which will have the follwing files in its lib directory: +```console +(venv) $ ls install/lib/ +cmake libggml.so libllama-common.so.0 libllama.so.0.1.0 llama.cpp +libggml-base.so libggml.so.0 libllama-common.so.0.1.0 libmtmd.so pkgconfig +libggml-base.so.0 libggml.so.0.19.0 libllama.so libmtmd.so.0 +libggml-base.so.0.19.0 libllama-common.so libllama.so.0 libmtmd.so.0.1.0 +``` + +Build/run this project using the installation created above: +```console +(venv) $ ./build.sh +-- Configuring done (0.0s) +-- Generating done (0.0s) +-- Build files have been written to: /home/danbev/work/ai/llama.cpp/examples/test-cmake/build +[100%] Built target test-cmake +[test-cmake] Using llama.cpp version 0.1.0-dev-b10335 +[test-cmake] Initializing backend... +load_backend: loaded CPU backend from /home/danbev/work/ai/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so +[test-cmake] Backend initialized. +``` diff --git a/examples/test-cmake/build-install.sh b/examples/test-cmake/build-install.sh new file mode 100755 index 000000000000..77a6713d67c4 --- /dev/null +++ b/examples/test-cmake/build-install.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +set -e + +rm -rf llama-build-install install + +cmake --fresh -S ../../. -B llama-build-install -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_SHARED_LIBS=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DLLAMA_TESTS_INSTALL=OFF \ + -DCMAKE_INSTALL_PREFIX="${PWD}/install" \ + -DGGML_BACKEND_DIR="${PWD}/install/lib/llama.cpp" \ + -DGGML_LIB_INSTALL_DIR="${PWD}/install/lib/llama.cpp" \ + -DLLAMA_LIB_INSTALL_DIR="${PWD}/install/lib/llama.cpp" \ + -DLLAMA_TOOLS_INSTALL=OFF + +cmake --build llama-build-install --parallel 12 +cmake --install llama-build-install diff --git a/examples/test-cmake/build.sh b/examples/test-cmake/build.sh new file mode 100755 index 000000000000..a212732b89d9 --- /dev/null +++ b/examples/test-cmake/build.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +set -e + +cmake -S . -B build -DCMAKE_PREFIX_PATH="${PWD}/install" +cmake --build build +LD_LIBRARY_PATH="${PWD}/install/lib/llama.cpp:${PWD}/install/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" ./build/test-cmake diff --git a/examples/test-cmake/test-cmake.cpp b/examples/test-cmake/test-cmake.cpp new file mode 100644 index 000000000000..c5c4765b439c --- /dev/null +++ b/examples/test-cmake/test-cmake.cpp @@ -0,0 +1,12 @@ +#include "llama.h" +#include + +int main(void) { + printf("[test-cmake] version: %s, build: %d (%s)\n", + llama_version(), LLAMA_BUILD_NUMBER, LLAMA_BUILD_COMMIT); + printf("[test-cmake] Initializing backend...\n"); + llama_backend_init(); + printf("[test-cmake] Backend initialized.\n"); + llama_backend_free(); + return 0; +} diff --git a/include/llama.h b/include/llama.h index ef278c9c3238..177fc10a9139 100644 --- a/include/llama.h +++ b/include/llama.h @@ -457,6 +457,8 @@ extern "C" { // lora adapter struct llama_adapter_lora; + LLAMA_API const char * llama_version(void); + // Helpers for getting default parameters // TODO: update API to start accepting pointers to params structs (https://github.com/ggml-org/llama.cpp/discussions/9172) LLAMA_API struct llama_model_params llama_model_default_params(void); diff --git a/scripts/make-release-checks.sh b/scripts/make-release-checks.sh new file mode 100755 index 000000000000..c8c6322841da --- /dev/null +++ b/scripts/make-release-checks.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# Run all pre-release checks and determine the release version. +# +# Usage: make-release-checks.sh [--dry-run] +# --dry-run: warn on failures instead of aborting +# +# Env (when running in GitHub Actions): GH_TOKEN, GITHUB_REPOSITORY, GITHUB_OUTPUT +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" + +DRY_RUN=false +for arg in "$@"; do + case "$arg" in + --dry-run) DRY_RUN=true ;; + *) echo "Unknown argument: $arg"; exit 1 ;; + esac +done + +MAJOR=$(grep "set(LLAMA_VERSION_MAJOR" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+') +MINOR=$(grep "set(LLAMA_VERSION_MINOR" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+') +PATCH=$(grep "set(LLAMA_VERSION_PATCH" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+') +VERSION="v${MAJOR}.${MINOR}.${PATCH}" +echo "Determined version: ${VERSION}" +if [[ -n "${GITHUB_OUTPUT:-}" ]]; then + echo "version=${VERSION}" >> "$GITHUB_OUTPUT" +fi + +echo "Checking that tag ${VERSION} does not already exist..." +if git ls-remote --tags origin "${VERSION}" | grep -q "${VERSION}"; then + echo "Error: tag ${VERSION} already exists on remote" + exit 1 +fi +echo "Tag ${VERSION} does not exist on remote - OK" + +SHA=$(git rev-parse HEAD) +echo "Checking release.yml status for commit ${SHA}..." +if [[ -z "${GITHUB_REPOSITORY:-}" ]]; then + echo "Warning: GITHUB_REPOSITORY not set - skipping CI check (local run)" +else + RUNS=$(gh api "repos/${GITHUB_REPOSITORY}/actions/workflows/release.yml/runs" \ + --jq "[.workflow_runs[] | select(.head_sha == \"${SHA}\" and .conclusion == \"success\")] | length") + if [[ "$RUNS" -eq 0 ]]; then + if [[ "$DRY_RUN" == "true" ]]; then + echo "Warning: no successful release.yml run found for HEAD (${SHA}) (dry run, continuing)." + else + echo "Error: no successful release.yml run found for HEAD (${SHA})" + echo "The nightly build must complete successfully before making a release." + exit 1 + fi + else + echo "Found successful release.yml run for HEAD." + fi +fi + +MAJOR=$(grep "set(GGML_VERSION_MAJOR" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+') +MINOR=$(grep "set(GGML_VERSION_MINOR" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+') +PATCH=$(grep "set(GGML_VERSION_PATCH" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+') +GGML_VERSION="v${MAJOR}.${MINOR}.${PATCH}" +echo "Local ggml version: ${GGML_VERSION}" + +if ! git clone --depth 1 --branch "${GGML_VERSION}" https://github.com/ggml-org/ggml.git upstream-ggml 2>/dev/null; then + echo "Warning: tag ${GGML_VERSION} not found in upstream ggml - skipping comparison" +else + echo "Comparing local ggml/ src and include with upstream ${GGML_VERSION}..." + DIFF=$(diff -rq "$REPO_ROOT/ggml/src" upstream-ggml/src 2>&1 || true) + DIFF+=$(diff -rq "$REPO_ROOT/ggml/include" upstream-ggml/include 2>&1 || true) + DIFF+=$(diff "$REPO_ROOT/ggml/CMakeLists.txt" upstream-ggml/CMakeLists.txt 2>&1 || true) + rm -rf upstream-ggml + if [[ -n "$DIFF" ]]; then + echo "local ggml/ differs from upstream ${GGML_VERSION}:" + echo "$DIFF" + if [[ "$DRY_RUN" == "true" ]]; then + echo "Warning: would abort release due to ggml mismatch (dry run, continuing)." + else + echo "Error: ggml must match upstream before making a release." + exit 1 + fi + else + echo "local ggml/ matches upstream ${GGML_VERSION}" + fi +fi diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 24f05cc91673..39ba3061f704 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -45,11 +45,16 @@ add_library(llama ) set_target_properties(llama PROPERTIES - VERSION ${LLAMA_INSTALL_VERSION} - SOVERSION 0 + VERSION ${LLAMA_VERSION_BASE} + SOVERSION ${LLAMA_VERSION_MAJOR} MACHO_CURRENT_VERSION 0 # keep macOS linker from seeing oversized version number ) +target_compile_definitions(llama PRIVATE + LLAMA_VERSION="${LLAMA_VERSION}" + LLAMA_COMMIT="${LLAMA_BUILD_COMMIT}" +) + target_include_directories(llama PRIVATE .) target_include_directories(llama PUBLIC ../include) target_compile_features (llama PRIVATE cxx_std_17) # don't bump diff --git a/src/llama.cpp b/src/llama.cpp index 94c8f60e0c43..9ff1902fc1d5 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -114,6 +114,10 @@ bool llama_supports_rpc(void) { return ggml_backend_reg_by_name("RPC") != nullptr; } +const char * llama_version(void) { + return LLAMA_VERSION; +} + void llama_backend_init(void) { ggml_time_init(); diff --git a/tests/test-quantize-stats.cpp b/tests/test-quantize-stats.cpp index c65557534025..e07d75b7e767 100644 --- a/tests/test-quantize-stats.cpp +++ b/tests/test-quantize-stats.cpp @@ -301,7 +301,7 @@ int main(int argc, char ** argv) { return 1; } - llama_print_build_info(); + llama_print_build_info(llama_version()); // load the model fprintf(stderr, "Loading model\n"); diff --git a/tools/cvector-generator/cvector-generator.cpp b/tools/cvector-generator/cvector-generator.cpp index 8c6b3d868d29..558c37e61298 100644 --- a/tools/cvector-generator/cvector-generator.cpp +++ b/tools/cvector-generator/cvector-generator.cpp @@ -421,7 +421,7 @@ int main(int argc, char ** argv) { params.cb_eval_user_data = &cb_data; params.warmup = false; - llama_print_build_info(); + llama_print_build_info(llama_version()); llama_backend_init(); llama_numa_init(params.numa); diff --git a/tools/gguf-split/gguf-split.cpp b/tools/gguf-split/gguf-split.cpp index 5cafcc9aa96a..c6cdbb98e278 100644 --- a/tools/gguf-split/gguf-split.cpp +++ b/tools/gguf-split/gguf-split.cpp @@ -106,7 +106,7 @@ static void split_params_parse_ex(int argc, const char ** argv, split_params & p split_print_usage(argv[0]); exit(0); } else if (arg == "--version") { - fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit()); + fprintf(stderr, "version: %s (build %d, commit %s)\n", llama_version(), llama_build_number(), llama_commit()); fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target()); exit(0); } else if (arg == "--dry-run") { diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 3f4a6c670dd7..769a44e0b73d 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -72,8 +72,8 @@ add_library(mtmd ) set_target_properties(mtmd PROPERTIES - VERSION ${LLAMA_INSTALL_VERSION} - SOVERSION 0 + VERSION ${LLAMA_VERSION_BASE} + SOVERSION ${LLAMA_VERSION_MAJOR} MACHO_CURRENT_VERSION 0 # keep macOS linker from seeing oversized version number ) diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 15ef64c4b0ed..8d03c8fcd427 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -611,7 +611,7 @@ int llama_quantize(int argc, char ** argv) { } } - llama_print_build_info(); + llama_print_build_info(llama_version()); if (params.dry_run) { fprintf(stderr, "%s: calculating quantization size for '%s' as %s", __func__, fname_inp.c_str(), ftype_str.c_str()); From 7a9ff95979d0d9a2fa79962c52f0e929cc054a7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Wed, 12 Aug 2026 15:41:43 +0200 Subject: [PATCH 040/211] disable rocm cache (#26962) --- .github/workflows/build-cache.yml | 48 ++++++++++++------------ .github/workflows/build-cuda-windows.yml | 14 +++---- .github/workflows/release.yml | 30 +++++++-------- 3 files changed, 46 insertions(+), 46 deletions(-) diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index e15fc5e0830f..2a103172850f 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -119,27 +119,27 @@ jobs: version_major: ${{ env.OPENVINO_VERSION_MAJOR }} version_full: ${{ env.OPENVINO_VERSION_FULL }} - windows-2022-rocm-cache: - runs-on: windows-2022 - - env: - # Make sure this is in sync with release.yml and build-cuda-windows.yml - ROCM_VERSION: "7.14.0" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Setup Cache - uses: actions/cache@v5 - id: cache-rocm - with: - path: C:\TheRock\build - key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - - - name: Setup ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' - uses: ./.github/actions/windows-setup-rocm - with: - version: ${{ env.ROCM_VERSION }} + # windows-2022-rocm-cache: + # runs-on: windows-2022 + + # env: + # # Make sure this is in sync with release.yml and build-cuda-windows.yml + # ROCM_VERSION: "7.14.0" + + # steps: + # - name: Clone + # id: checkout + # uses: actions/checkout@v6 + + # - name: Setup Cache + # uses: actions/cache@v5 + # id: cache-rocm + # with: + # path: C:\TheRock\build + # key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} + + # - name: Setup ROCm + # if: steps.cache-rocm.outputs.cache-hit != 'true' + # uses: ./.github/actions/windows-setup-rocm + # with: + # version: ${{ env.ROCM_VERSION }} diff --git a/.github/workflows/build-cuda-windows.yml b/.github/workflows/build-cuda-windows.yml index ff900802f4a2..8b59f3975c5c 100644 --- a/.github/workflows/build-cuda-windows.yml +++ b/.github/workflows/build-cuda-windows.yml @@ -97,15 +97,15 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Cache ROCm Installation - uses: actions/cache@v5 - id: cache-rocm - with: - path: C:\TheRock\build - key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} + # - name: Cache ROCm Installation + # uses: actions/cache@v5 + # id: cache-rocm + # with: + # path: C:\TheRock\build + # key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' + # if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: version: ${{ env.ROCM_VERSION }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index af34c9499396..ad0df09f3860 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -774,15 +774,15 @@ jobs: key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} evict-old-files: 1d - - name: Cache ROCm Installation - id: cache-rocm - uses: actions/cache@v5 - with: - path: C:\TheRock\build - key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }} + # - name: Cache ROCm Installation + # id: cache-rocm + # uses: actions/cache@v5 + # with: + # path: C:\TheRock\build + # key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' + # if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: version: ${{ matrix.ROCM_VERSION }} @@ -1320,10 +1320,10 @@ jobs: with: tool-cache: true - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - with: - key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} + # - name: ccache + # uses: ggml-org/ccache-action@v1.2.21 + # with: + # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} - name: Dependencies id: depends @@ -1379,10 +1379,10 @@ jobs: ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(nproc) - - name: ccache-clear - uses: ./.github/actions/ccache-clear - with: - key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} + # - name: ccache-clear + # uses: ./.github/actions/ccache-clear + # with: + # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} - name: Determine tag name id: tag From 9558fa44c92746a58dd07ad1bf0c889715b938a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Wed, 12 Aug 2026 15:41:44 +0200 Subject: [PATCH 041/211] ci : disable ubuntu-rocm (#26969) * disable ubuntu-rocm * link PR --- .github/workflows/release.yml | 238 +++++++++++++++++----------------- 1 file changed, 119 insertions(+), 119 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ad0df09f3860..82a8364e435c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1285,123 +1285,123 @@ jobs: path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-${{ matrix.build }}-x64.tar.gz name: llama-bin-ubuntu-sycl-${{ matrix.build }}-x64.tar.gz - ubuntu-22-rocm: - needs: [check-release, get-version] - if: ${{ needs.check-release.outputs.should_release == 'true' }} - - runs-on: ubuntu-22.04 - - permissions: - actions: write - - strategy: - matrix: - include: - - ROCM_VERSION: "7.14.0" - gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201" - build: 'x64' - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - - name: Free up disk space - uses: ggml-org/free-disk-space@v1.3.1 - with: - tool-cache: true - - # - name: ccache - # uses: ggml-org/ccache-action@v1.2.21 - # with: - # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} - - - name: Dependencies - id: depends - run: | - sudo apt install -y build-essential git cmake wget - - - name: Setup TheRock with Wheels - id: therock_env - run: | - # Create Python virtual environment - python3 -m venv .venv - source .venv/bin/activate - - # Install ROCm wheels for build - # libraries = HIP runtime and CMake configs needed for linking - # devel = compilers, headers, static libs - python -m pip install --upgrade pip - python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}" - - # Get ROCm installation paths using the rocm-sdk CLI tool - ROCM_PATH=$(rocm-sdk path --root) - CMAKE_PATH=$(rocm-sdk path --cmake) - BIN_PATH=$(rocm-sdk path --bin) - echo "ROCM_PATH=$ROCM_PATH" - echo "CMAKE_PATH=$CMAKE_PATH" - echo "BIN_PATH=$BIN_PATH" - - # Set environment variables - echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV - echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV - echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV - echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV - echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV - - # Keep venv activated for subsequent steps - echo "$(pwd)/.venv/bin" >> $GITHUB_PATH - - - name: Build with native CMake HIP support - id: cmake_build - run: | - cmake -B build -S . \ - -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \ - -DCMAKE_BUILD_TYPE=Release \ - -DGGML_BACKEND_DL=ON \ - -DGGML_NATIVE=OFF \ - -DCMAKE_INSTALL_RPATH='$ORIGIN' \ - -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ - -DGGML_CPU_ALL_VARIANTS=ON \ - -DGPU_TARGETS="${{ matrix.gpu_targets }}" \ - -DGGML_HIP=ON \ - -DHIP_PLATFORM=amd \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ - ${{ env.CMAKE_ARGS }} - cmake --build build --config Release -j $(nproc) - - # - name: ccache-clear - # uses: ./.github/actions/ccache-clear - # with: - # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} - - - name: Determine tag name - id: tag - uses: ./.github/actions/get-tag-name - - - name: Get ROCm short version - run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV - - - name: Pack artifacts - id: pack_artifacts - run: | - cp LICENSE ./build/bin/ - tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin . - - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz - name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz +# ubuntu-22-rocm: +# needs: [check-release, get-version] +# if: ${{ needs.check-release.outputs.should_release == 'true' }} + +# runs-on: ubuntu-22.04 + +# permissions: +# actions: write + +# strategy: +# matrix: +# include: +# - ROCM_VERSION: "7.14.0" +# gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201" +# build: 'x64' + +# steps: +# - name: Clone +# id: checkout +# uses: actions/checkout@v6 +# with: +# fetch-depth: 0 + +# - name: Setup Node.js +# uses: actions/setup-node@v6 +# with: +# node-version: "24" +# cache: "npm" +# cache-dependency-path: "tools/ui/package-lock.json" + +# - name: Free up disk space +# uses: ggml-org/free-disk-space@v1.3.1 +# with: +# tool-cache: true + +# # - name: ccache +# # uses: ggml-org/ccache-action@v1.2.21 +# # with: +# # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} + +# - name: Dependencies +# id: depends +# run: | +# sudo apt install -y build-essential git cmake wget + +# - name: Setup TheRock with Wheels +# id: therock_env +# run: | +# # Create Python virtual environment +# python3 -m venv .venv +# source .venv/bin/activate + +# # Install ROCm wheels for build +# # libraries = HIP runtime and CMake configs needed for linking +# # devel = compilers, headers, static libs +# python -m pip install --upgrade pip +# python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}" + +# # Get ROCm installation paths using the rocm-sdk CLI tool +# ROCM_PATH=$(rocm-sdk path --root) +# CMAKE_PATH=$(rocm-sdk path --cmake) +# BIN_PATH=$(rocm-sdk path --bin) +# echo "ROCM_PATH=$ROCM_PATH" +# echo "CMAKE_PATH=$CMAKE_PATH" +# echo "BIN_PATH=$BIN_PATH" + +# # Set environment variables +# echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV +# echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV +# echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV +# echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV +# echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV + +# # Keep venv activated for subsequent steps +# echo "$(pwd)/.venv/bin" >> $GITHUB_PATH + +# - name: Build with native CMake HIP support +# id: cmake_build +# run: | +# cmake -B build -S . \ +# -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \ +# -DCMAKE_BUILD_TYPE=Release \ +# -DGGML_BACKEND_DL=ON \ +# -DGGML_NATIVE=OFF \ +# -DCMAKE_INSTALL_RPATH='$ORIGIN' \ +# -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ +# -DGGML_CPU_ALL_VARIANTS=ON \ +# -DGPU_TARGETS="${{ matrix.gpu_targets }}" \ +# -DGGML_HIP=ON \ +# -DHIP_PLATFORM=amd \ +# -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ +# ${{ env.CMAKE_ARGS }} +# cmake --build build --config Release -j $(nproc) + +# # - name: ccache-clear +# # uses: ./.github/actions/ccache-clear +# # with: +# # key: release-ubuntu-22.04-rocm-${{ matrix.ROCM_VERSION }} + +# - name: Determine tag name +# id: tag +# uses: ./.github/actions/get-tag-name + +# - name: Get ROCm short version +# run: echo "ROCM_VERSION_SHORT=$(echo '${{ matrix.ROCM_VERSION }}' | cut -d '.' -f 1,2)" >> $GITHUB_ENV + +# - name: Pack artifacts +# id: pack_artifacts +# run: | +# cp LICENSE ./build/bin/ +# tar -czvf llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C ./build/bin . + +# - name: Upload artifacts +# uses: actions/upload-artifact@v6 +# with: +# path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz +# name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz ios-xcode: needs: [check-release, get-version] @@ -1578,7 +1578,7 @@ jobs: #- windows-sycl - windows-rocm - windows-openvino - - ubuntu-22-rocm + #- ubuntu-22-rocm - ubuntu-cpu - ubuntu-vulkan - ubuntu-24-openvino @@ -1688,7 +1688,7 @@ jobs: - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz) - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz) - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz) - - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz) + - Ubuntu x64 (ROCm 7.14)[DISABLED](https://github.com/ggml-org/llama.cpp/pull/26969) - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz) From 84e908c625fb60992b4cdef8180fb12fa9b4c4bf Mon Sep 17 00:00:00 2001 From: Eve <139727413+netrunnereve@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:01:12 +0000 Subject: [PATCH 042/211] ci: fix thread sanitizer + remove ccache (#26927) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test address on Intel-LNL-U7-258V * retry * run address on github * use native build for cpu * this should be runnable everywhere multicore * disable ccache --------- Co-authored-by: Sigbjørn Skjæret --- .github/workflows/build-cmake-pkg.yml | 6 +++--- .github/workflows/build-cpu.yml | 1 + .github/workflows/build-sanitize.yml | 20 ++++++++++---------- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/.github/workflows/build-cmake-pkg.yml b/.github/workflows/build-cmake-pkg.yml index a958870de3d6..0e4069ce3526 100644 --- a/.github/workflows/build-cmake-pkg.yml +++ b/.github/workflows/build-cmake-pkg.yml @@ -5,7 +5,7 @@ on: jobs: linux: - runs-on: [self-hosted, Linux, CPU] + runs-on: [self-hosted, Linux] steps: - uses: actions/checkout@v6 with: @@ -23,7 +23,7 @@ jobs: -DLLAMA_BUILD_APP=OFF \ -DLLAMA_BUILD_IS_DEV=OFF \ -DCMAKE_BUILD_TYPE=Release - cmake --build build --config Release + cmake --build build --config Release -j $(nproc) cmake --install build --prefix "$PREFIX" --config Release export LLAMA_CONFIG="$PREFIX"/lib/cmake/llama/llama-config.cmake @@ -54,4 +54,4 @@ jobs: cd examples/simple-cmake-pkg cmake -S . -B build -DCMAKE_PREFIX_PATH="$PREFIX"/lib/cmake - cmake --build build + cmake --build build -j $(nproc) diff --git a/.github/workflows/build-cpu.yml b/.github/workflows/build-cpu.yml index 57622e1d4507..df70f4d10bc6 100644 --- a/.github/workflows/build-cpu.yml +++ b/.github/workflows/build-cpu.yml @@ -94,6 +94,7 @@ jobs: id: cmake_build run: | cmake -B build \ + -DGGML_NATIVE=OFF \ -DLLAMA_FATAL_WARNINGS=ON \ -DGGML_RPC=ON \ -DGGML_NATIVE=OFF diff --git a/.github/workflows/build-sanitize.yml b/.github/workflows/build-sanitize.yml index 9654cb4e4d91..974af62eb2e6 100644 --- a/.github/workflows/build-sanitize.yml +++ b/.github/workflows/build-sanitize.yml @@ -39,9 +39,9 @@ jobs: strategy: matrix: include: + # thread and address doesn't run properly on some self hosted machines, so run it on Github instead - sanitizer: ADDRESS - machine: [self-hosted, X64, Linux] - # thread doesn't run properly on some self hosted machines, so run it on Github instead + machine: ubuntu-24.04 - sanitizer: THREAD machine: ubuntu-24.04 - sanitizer: UNDEFINED @@ -54,14 +54,14 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - if: ${{ matrix.sanitizer == 'THREAD' }} - with: - key: ctest-thread-ubuntu-24.04 - variant: ccache - evict-old-files: 1d - save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + # - name: ccache + # uses: ggml-org/ccache-action@v1.2.21 + # if: ${{ matrix.sanitizer != 'UNDEFINED' }} + # with: + # key: ctest-${{ matrix.sanitizer }}-ubuntu-24.04 + # variant: ccache + # evict-old-files: 1d + # save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} # with UNDEFINED sanitizer, we have to build in Debug to avoid GCC 13 false-positive warnings - name: Build (undefined) From 8e7f22b67ef4667b4ddd50230771287f328cfb3f Mon Sep 17 00:00:00 2001 From: Johnathan Craig Maudlin <13183098+jcmdln@users.noreply.github.com> Date: Wed, 12 Aug 2026 18:02:27 -0400 Subject: [PATCH 043/211] common: add system-level config file (#26118) * common: Add CLI > ENV > models-presets > INI precedence 1. CLI flags have the highest precedence 2. ENV vars have the second-highest precedence 3. System and User configs have the lowest precedence - Linux/BSD/Mac - /etc/llama.cpp/config.ini < ${XDG_CONFIG_HOME:-~/.config}/llama.cpp/config.ini - Windows - %PROGRAMDATA%\llama.cpp\config.ini < %APPDATA%\llama.cpp\config.ini * fix UB * use common_get_env * ignore_unknown_keys * nits * add docs --------- Co-authored-by: Xuan Son Nguyen --- common/arg.cpp | 50 ++++++++++++++++++++++++++++++++ common/common.cpp | 73 ++++++++++++++++++++++++++++++++++++++++------- common/common.h | 1 + common/preset.cpp | 2 ++ common/preset.h | 4 +++ docs/preset.md | 17 ++++++++++- 6 files changed, 136 insertions(+), 11 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index b2fddfe5f2f0..9fe7d66e2426 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include // for hardware_concurrency #include @@ -704,12 +705,61 @@ void common_models_handler_apply(common_models_handler & handler, common_params // CLI argument parsing functions // +// apply config files (if present), a later file overrides an earlier one: +// 1. system-wide: /etc/llama.cpp/config.ini (%PROGRAMDATA%\llama.cpp\config.ini on windows) +// 2. user-level: ${XDG_CONFIG_HOME:-~/.config}/llama.cpp/config.ini (%APPDATA%\llama.cpp\config.ini on windows) +static void common_params_apply_system_config(common_params & params, llama_example ex) { + std::vector paths; + +#if defined(_WIN32) + const std::string program_data = common_get_env("PROGRAMDATA"); + if (!program_data.empty()) { + paths.push_back(program_data + "\\llama.cpp\\config.ini"); + } +#else + paths.push_back("/etc/llama.cpp/config.ini"); +#endif + + try { + paths.push_back(fs_get_config_directory() + "config.ini"); + } catch (const std::exception & e) { + LOG_DBG("cannot read user-level config file, skipping: %s\n", e.what()); + } + + std::vector found; + for (const auto & path : paths) { + std::error_code ec; + if (std::filesystem::exists(path, ec)) { + found.push_back(path); + } + } + if (found.empty()) { + return; + } + + common_preset_context ctx(ex); + ctx.ignore_unknown_keys = true; // the same config file is shared by all programs + for (const auto & path : found) { + LOG_INF("using config file: %s\n", path.c_str()); + common_preset global; + common_presets presets = ctx.load_from_ini(path, global); + global.apply_to_params(params); + auto it = presets.find(COMMON_PRESET_DEFAULT_NAME); + if (it != presets.end()) { + it->second.apply_to_params(params); + } + } +} + static bool common_params_parse_ex(int argc, char ** argv, common_params_context & ctx_arg) { common_params & params = ctx_arg.params; // setup log directly from params.verbosity: see tools/cli/cli.cpp common_log_set_verbosity_thold(params.verbosity); + // config file applies first, so env variables and CLI arguments override it + common_params_apply_system_config(params, ctx_arg.ex); + std::unordered_map> arg_to_options; for (auto & opt : ctx_arg.options) { for (const auto & arg : opt.args) { diff --git a/common/common.cpp b/common/common.cpp index 2e3f14cd1c43..5639bdf412e3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1019,20 +1019,21 @@ std::string fs_get_cache_directory() { std::string cache_directory = ""; auto ensure_trailing_slash = [](std::string p) { // Make sure to add trailing slash - if (p.back() != DIRECTORY_SEPARATOR) { + if (p.empty() || p.back() != DIRECTORY_SEPARATOR) { p += DIRECTORY_SEPARATOR; } return p; }; - if (getenv("LLAMA_CACHE")) { - cache_directory = std::getenv("LLAMA_CACHE"); - } else { + cache_directory = common_get_env("LLAMA_CACHE"); + if (cache_directory.empty()) { #if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \ defined(__OpenBSD__) || defined(__NetBSD__) - if (std::getenv("XDG_CACHE_HOME")) { - cache_directory = std::getenv("XDG_CACHE_HOME"); - } else if (std::getenv("HOME")) { - cache_directory = std::getenv("HOME") + std::string("/.cache/"); + const std::string xdg_cache_home = common_get_env("XDG_CACHE_HOME"); + const std::string home = common_get_env("HOME"); + if (!xdg_cache_home.empty()) { + cache_directory = xdg_cache_home; + } else if (!home.empty()) { + cache_directory = home + "/.cache/"; } else { #if defined(__linux__) /* no $HOME is defined, fallback to getpwuid */ @@ -1047,9 +1048,16 @@ std::string fs_get_cache_directory() { #endif /* defined(__linux__) */ } #elif defined(__APPLE__) - cache_directory = std::getenv("HOME") + std::string("/Library/Caches/"); + cache_directory = common_get_env("HOME"); + if (cache_directory.empty()) { + throw std::runtime_error("Failed to find $HOME directory"); + } + cache_directory += "/Library/Caches/"; #elif defined(_WIN32) - cache_directory = std::getenv("LOCALAPPDATA"); + cache_directory = common_get_env("LOCALAPPDATA"); + if (cache_directory.empty()) { + throw std::runtime_error("Failed to find %LOCALAPPDATA% directory"); + } #elif defined(__EMSCRIPTEN__) GGML_ABORT("not implemented on this platform"); #else @@ -1061,6 +1069,51 @@ std::string fs_get_cache_directory() { return ensure_trailing_slash(cache_directory); } +std::string fs_get_config_directory() { + std::string config_directory = ""; + auto ensure_trailing_slash = [](std::string p) { + if (p.empty() || p.back() != DIRECTORY_SEPARATOR) { + p += DIRECTORY_SEPARATOR; + } + return p; + }; +#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \ + defined(__OpenBSD__) || defined(__NetBSD__) || defined(__APPLE__) + const std::string xdg_config_home = common_get_env("XDG_CONFIG_HOME"); + const std::string home = common_get_env("HOME"); + if (!xdg_config_home.empty()) { + config_directory = xdg_config_home; + } else if (!home.empty()) { + config_directory = home + "/.config/"; + } else { +#if defined(__linux__) + /* no $HOME is defined, fallback to getpwuid */ + struct passwd *pw = getpwuid(getuid()); + if ((!pw) || (!pw->pw_dir)) { + throw std::runtime_error("Failed to find $HOME directory"); + } + + config_directory = std::string(pw->pw_dir) + std::string("/.config/"); +#else + throw std::runtime_error("Failed to find $HOME directory"); +#endif + } +#elif defined(_WIN32) + config_directory = common_get_env("APPDATA"); + if (config_directory.empty()) { + throw std::runtime_error("Failed to find %APPDATA% directory"); + } +#elif defined(__EMSCRIPTEN__) + // caller decides what to do when there is no config directory + throw std::runtime_error("not implemented on this platform"); +#else +# error Unknown architecture +#endif + config_directory = ensure_trailing_slash(config_directory); + config_directory += "llama.cpp"; + return ensure_trailing_slash(config_directory); +} + std::string fs_get_cache_file(const std::string & filename) { GGML_ASSERT(filename.find(DIRECTORY_SEPARATOR) == std::string::npos); std::string cache_directory = fs_get_cache_directory(); diff --git a/common/common.h b/common/common.h index d485d4fb41f7..b16bd3bb6855 100644 --- a/common/common.h +++ b/common/common.h @@ -881,6 +881,7 @@ bool fs_is_directory(const std::string & path); std::string fs_get_cache_directory(); std::string fs_get_cache_file(const std::string & filename); +std::string fs_get_config_directory(); struct common_file_info { std::string path; diff --git a/common/preset.cpp b/common/preset.cpp index eb0c60b09cff..0b29af883426 100644 --- a/common/preset.cpp +++ b/common/preset.cpp @@ -322,6 +322,8 @@ common_presets common_preset_context::load_from_ini(const std::string & path, co preset.options[opt] = value; } LOG_DBG("accepted option: %s = %s\n", key.c_str(), preset.options[opt].c_str()); + } else if (ignore_unknown_keys) { + LOG_WRN("ignoring option '%s' from %s: not supported by this program\n", key.c_str(), path.c_str()); } else { throw std::runtime_error(string_format( "option '%s' not recognized in preset '%s'", diff --git a/common/preset.h b/common/preset.h index 52935ebde86e..d8fc3915bc8b 100644 --- a/common/preset.h +++ b/common/preset.h @@ -59,6 +59,10 @@ struct common_preset_context { bool filter_allowed_keys = false; std::set allowed_keys; + // if true, options unknown to the current example are skipped instead of being an error + // used for config files shared by all binaries, where each binary only knows a subset of options + bool ignore_unknown_keys = false; + // if only_remote_allowed is true, only accept whitelisted keys common_preset_context(llama_example ex); diff --git a/docs/preset.md b/docs/preset.md index 85762a420b31..3d85467e8700 100644 --- a/docs/preset.md +++ b/docs/preset.md @@ -4,7 +4,7 @@ The INI preset feature, introduced in [PR#17859](https://github.com/ggml-org/llama.cpp/pull/17859), allows users to create reusable and shareable parameter configurations for llama.cpp. -### Using Presets with the Server +## Using Presets with the Server When running multiple models on the server (router mode), INI preset files can be used to configure model-specific parameters. Please refer to the [server documentation](../tools/server/README.md) for more details. @@ -93,3 +93,18 @@ llama-server -hf user/repo:gpt-oss-120b-hf ``` Please make sure to provide the correct `hf-repo` for each child preset. Otherwise, you may get error: `The specified tag is not a valid quantization scheme.` + +## System-level config + +The system-level config, added in PR [#26118](https://github.com/ggml-org/llama.cpp/pull/26118), allows sharing the same set of options among multiple tools and examples. Unlike the sections above, it is not limited to the server. + +These files are loaded on startup if present. A later file overrides an earlier one: +1. System-wide: `/etc/llama.cpp/config.ini` (or `%PROGRAMDATA%\llama.cpp\config.ini` on Windows) +2. User-level: `$XDG_CONFIG_HOME/llama.cpp/config.ini`, `~/.config/llama.cpp/config.ini` by default (or `%APPDATA%\llama.cpp\config.ini` on Windows) + +The config file is applied first, then its options are overridden by ENV variables, CLI arguments and model presets (in router mode). + +Note: +- Only the `[*]` and default sections are used; options written before any section header belong to "default. Named sections are ignored +- Tool-specific options can be specified, but will be ignored (with a warning) if the example doesn't support it
Example: if you specify `port = 1234`, only `llama-server` will use it, other examples will ignore it +- `model` or `hf-repo` are not recommended to be configured system-level, because it may introduce conflicts
Example: a `hf-repo` in the config file still takes effect when you pass `-m` on the command line, so you may load a different model than expected From e21152dc9636c6d2db6edc9b4531dfc5b2d1cba3 Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Thu, 13 Aug 2026 06:53:56 +0200 Subject: [PATCH 044/211] ui: Constants refactor (#26908) * refactor: Constants * refactor: Constants/Enums cleanup * refactor: Constant objects instead of multiple single value constants * refactor: Cleanup constants --- tools/ui/pwa-assets-dark.config.ts | 2 +- tools/ui/pwa-assets.config.ts | 2 +- tools/ui/scripts/vite-plugin-build-info.ts | 2 +- .../ui/scripts/vite-plugin-relativize-base.ts | 2 +- tools/ui/scripts/vite-plugin-splash-screen.ts | 9 +- .../actions/ActionIconCopyToClipboard.svelte | 2 +- ...hatAttachmentsListItemThumbnailFile.svelte | 2 +- ...hatAttachmentsPreviewCurrentItemPdf.svelte | 2 +- ...hatAttachmentsPreviewThumbnailStrip.svelte | 2 +- .../ChatFormActionAddButton.svelte | 3 +- .../ChatFormActionAddDropdown.svelte | 2 +- .../ChatFormActionAddMcpServersSubmenu.svelte | 3 +- .../ChatFormActionAddReasoningSubmenu.svelte | 2 +- .../ChatFormActionAddSheet.svelte | 8 +- .../ChatFormActionAddToolsSubmenu.svelte | 3 +- .../ChatFormActionRecord.svelte | 2 +- .../ChatFormActions/ChatFormActions.svelte | 3 +- .../ChatForm/ChatFormContenteditable.svelte | 6 +- .../ChatFormMentionPicker.svelte | 10 +- .../ChatForm/ChatFormWorkingDirectory.svelte | 23 +--- .../ChatMessage/ChatMessage.svelte | 4 +- .../ChatMessageToolCallBlockDefault.svelte | 2 +- ...essageToolCallBlockExecShellCommand.svelte | 3 +- .../ChatMessageToolCallBlockReadFile.svelte | 4 +- .../ChatMessageToolCallBlockReadMedia.svelte | 2 +- ...atMessageToolCallBlockSearchResults.svelte | 2 +- .../ChatMessageToolCall/ToolCallBlock.svelte | 5 +- .../ChatMessageToolCall/parsers/read-file.ts | 10 +- .../ChatMessageToolCall/parsers/read-media.ts | 2 +- .../ChatMessageToolCall/parsers/write-file.ts | 9 +- .../ChatMessageActionCard.svelte | 2 +- .../ChatMessageReasoningBlock.svelte | 2 +- .../ChatScreenActionScrollDown.svelte | 2 +- .../ChatScreen/ChatScreenServerError.svelte | 2 +- .../content/CollapsibleContentBlock.svelte | 2 +- .../MarkdownContent/MarkdownContent.svelte | 49 +++---- .../plugins/rehype/code-block-utils.ts | 22 ++-- .../plugins/rehype/enhance-code-blocks.ts | 6 +- .../plugins/rehype/enhance-svg-blocks.ts | 35 ++--- .../plugins/rehype/file-badge.ts | 2 +- .../MarkdownContent/plugins/rehype/svg-pre.ts | 8 +- .../app/content/MermaidPreview.svelte | 4 +- .../app/content/MermaidPreviewControls.svelte | 2 +- .../app/content/SyntaxHighlightedCode.svelte | 2 +- .../dialogs/DialogMcpResourcePreview.svelte | 12 +- .../dialogs/DialogMcpResourcesBrowser.svelte | 2 +- .../app/dialogs/DialogMcpServerAddNew.svelte | 9 +- .../dialogs/DialogModelNotAvailable.svelte | 3 +- .../components/app/forms/SearchInput.svelte | 2 +- .../app/mcp/McpActiveServersAvatars.svelte | 3 +- .../app/mcp/McpResourcePreview.svelte | 2 +- .../McpResourcesBrowserHeader.svelte | 2 +- .../McpResourcesBrowserServerItem.svelte | 2 +- .../mcp/McpServerCard/McpServerCard.svelte | 2 +- .../components/app/mcp/McpServerForm.svelte | 16 +-- .../app/misc/HorizontalScrollCarousel.svelte | 2 +- .../app/models/ModelsSelectorOption.svelte | 2 +- .../SidebarNavigationActions.svelte | 2 +- .../SidebarNavigationConversationItem.svelte | 3 +- .../app/server/ServerErrorSplash.svelte | 8 +- .../components/app/server/ServerStatus.svelte | 2 +- .../SettingsChat/SettingsChatFields.svelte | 3 +- .../SettingsChatImportExportSection.svelte | 2 +- .../SettingsChat/SettingsChatToolsTab.svelte | 4 +- .../SettingsChatDesktopSidebar.svelte | 2 +- .../settings/SettingsChatMobileHeader.svelte | 2 +- .../src/lib/components/pwa/PwaMetaTags.svelte | 3 +- .../{agentic.ts => agentic.constants.ts} | 16 +-- ...ndpoints.ts => api-endpoints.constants.ts} | 0 .../constants/{app.ts => app.constants.ts} | 0 .../ui/src/lib/constants/attachment-labels.ts | 4 - ...t-menu.ts => attachment-menu.constants.ts} | 25 +--- ...uto-scroll.ts => auto-scroll.constants.ts} | 0 ...ction.ts => binary-detection.constants.ts} | 0 ...n-tools.ts => built-in-tools.constants.ts} | 14 +- tools/ui/src/lib/constants/cache.constants.ts | 44 +++++++ tools/ui/src/lib/constants/cache.ts | 54 -------- .../{chat-form.ts => chat-form.constants.ts} | 0 .../{cli-flags.ts => cli-flags.constants.ts} | 0 .../src/lib/constants/code-block.constants.ts | 50 ++++++++ tools/ui/src/lib/constants/code-blocks.ts | 8 -- tools/ui/src/lib/constants/code.ts | 27 ---- .../constants/content-detection.constants.ts | 20 +++ ...up.ts => context-gauge-popup.constants.ts} | 0 ...text-keys.ts => context-keys.constants.ts} | 0 ...ctions.ts => control-actions.constants.ts} | 2 - ...rt.ts => conversation-import.constants.ts} | 0 ...ss-classes.ts => css-classes.constants.ts} | 0 .../{database.ts => database.constants.ts} | 2 +- ...-blocks.ts => diagram-blocks.constants.ts} | 0 .../{error.ts => error.constants.ts} | 0 .../lib/constants/floating-ui-constraints.ts | 2 - ...{formatters.ts => formatters.constants.ts} | 0 .../ui/src/lib/constants/headers.constants.ts | 35 +++++ .../{icons.ts => icons.constants.ts} | 0 tools/ui/src/lib/constants/image-size.ts | 3 - tools/ui/src/lib/constants/image.constants.ts | 32 +++++ tools/ui/src/lib/constants/index.ts | 120 +++++++++--------- tools/ui/src/lib/constants/jpeg-exif.ts | 30 ----- ...-pairs.ts => key-value-pairs.constants.ts} | 0 ...ction.ts => latex-protection.constants.ts} | 0 ...eral-html.ts => literal-html.constants.ts} | 5 - .../src/lib/constants/markdown.constants.ts | 17 +++ tools/ui/src/lib/constants/markdown.ts | 5 - ...e-size.ts => max-bundle-size.constants.ts} | 0 .../{mcp-form.ts => mcp-form.constants.ts} | 0 ...-resource.ts => mcp-resource.constants.ts} | 0 .../constants/{mcp.ts => mcp.constants.ts} | 64 +++------- ...on-badge.ts => mention-badge.constants.ts} | 0 ...-blocks.ts => mermaid-blocks.constants.ts} | 0 .../lib/constants/message-export.constants.ts | 24 ++++ tools/ui/src/lib/constants/message-export.ts | 23 ---- .../src/lib/constants/model-id.constants.ts | 43 +++++++ tools/ui/src/lib/constants/model-id.ts | 46 ------- ...-loading.ts => model-loading.constants.ts} | 0 ...h-display.ts => path-display.constants.ts} | 0 .../{precision.ts => precision.constants.ts} | 0 tools/ui/src/lib/constants/processing-info.ts | 8 -- .../constants/{pwa.ts => pwa.constants.ts} | 2 +- .../lib/constants/reasoning-effort-tokens.ts | 12 -- ...ffort.ts => reasoning-effort.constants.ts} | 11 ++ ...s => recommended-mcp-servers.constants.ts} | 0 .../{routes.ts => routes.constants.ts} | 0 .../ui/src/lib/constants/sandbox.constants.ts | 13 ++ ...ngs-keys.ts => settings-keys.constants.ts} | 0 ...stry.ts => settings-registry.constants.ts} | 24 ++-- .../constants/special-characters.constants.ts | 16 +++ .../{storage.ts => storage.constants.ts} | 0 .../constants/{sse.ts => stream.constants.ts} | 8 ++ tools/ui/src/lib/constants/stream.ts | 3 - ...s.ts => supported-file-types.constants.ts} | 0 .../src/lib/constants/svg-blocks.constants.ts | 57 +++++++++ tools/ui/src/lib/constants/svg-blocks.ts | 49 ------- ...er.ts => table-html-restorer.constants.ts} | 0 ...ation.ts => title-generation.constants.ts} | 0 tools/ui/src/lib/constants/tools.ts | 22 ---- tools/ui/src/lib/constants/tooltip-config.ts | 1 - .../lib/constants/{ui.ts => ui.constants.ts} | 43 +++++-- ...-template.ts => uri-template.constants.ts} | 26 +--- .../constants/{url.ts => url.constants.ts} | 0 tools/ui/src/lib/constants/viewport.ts | 1 - .../constants/working-directory.constants.ts | 50 ++++++++ .../ui/src/lib/constants/working-directory.ts | 49 ------- tools/ui/src/lib/enums/attachment.enums.ts | 10 ++ tools/ui/src/lib/enums/index.ts | 7 +- .../lib/hooks/use-chat-form-pickers.svelte.ts | 3 +- .../hooks/use-keyboard-shortcuts.svelte.ts | 2 +- tools/ui/src/lib/hooks/use-pwa.svelte.ts | 3 +- .../lib/hooks/use-reasoning-menu.svelte.ts | 3 +- .../hooks/use-settings-navigation.svelte.ts | 2 +- tools/ui/src/lib/services/chat.service.ts | 21 +-- tools/ui/src/lib/services/index.ts | 2 +- tools/ui/src/lib/services/mcp.service.ts | 14 +- .../ui/src/lib/services/migration.service.ts | 5 +- tools/ui/src/lib/services/models.service.ts | 50 +++----- tools/ui/src/lib/services/router.service.ts | 2 +- tools/ui/src/lib/services/tools.service.ts | 10 +- tools/ui/src/lib/stores/agentic.svelte.ts | 5 +- tools/ui/src/lib/stores/chat.svelte.ts | 20 +-- .../ui/src/lib/stores/conversations.svelte.ts | 32 ++--- .../ui/src/lib/stores/mcp-resources.svelte.ts | 7 +- tools/ui/src/lib/stores/mcp.svelte.ts | 17 +-- tools/ui/src/lib/stores/models.svelte.ts | 7 +- tools/ui/src/lib/stores/tools.svelte.ts | 2 +- tools/ui/src/lib/stores/viewport.svelte.ts | 2 +- tools/ui/src/lib/types/chat.d.ts | 47 ++++++- tools/ui/src/lib/types/index.ts | 10 +- tools/ui/src/lib/types/navigation.d.ts | 14 ++ tools/ui/src/lib/types/tools.d.ts | 10 ++ tools/ui/src/lib/utils/agentic.ts | 34 ++--- tools/ui/src/lib/utils/api-fetch.ts | 2 +- tools/ui/src/lib/utils/api-headers.ts | 22 ++-- tools/ui/src/lib/utils/api-key-validation.ts | 6 +- tools/ui/src/lib/utils/built-in-tools.ts | 13 ++ tools/ui/src/lib/utils/cache-ttl.ts | 10 +- tools/ui/src/lib/utils/cap-img-size.ts | 5 +- .../lib/{constants => utils}/chat-commands.ts | 17 +-- tools/ui/src/lib/utils/code.ts | 27 ++-- .../lib/utils/contenteditable-tokenizer.ts | 2 +- tools/ui/src/lib/utils/cors-proxy.ts | 10 +- tools/ui/src/lib/utils/glob-search.ts | 13 +- tools/ui/src/lib/utils/heic-to-jpeg.ts | 4 +- tools/ui/src/lib/utils/index.ts | 11 +- tools/ui/src/lib/utils/jpeg-orientation.ts | 32 ++--- tools/ui/src/lib/utils/mcp.ts | 47 +++---- tools/ui/src/lib/utils/mention-badge.ts | 7 +- tools/ui/src/lib/utils/path-display.ts | 6 +- .../sandbox.ts => utils/sandbox-tool.ts} | 19 +-- tools/ui/src/lib/utils/sanitize-svg.ts | 10 +- tools/ui/src/lib/utils/stream-identity.ts | 4 +- tools/ui/src/lib/utils/uri-template.ts | 42 +++--- tools/ui/src/lib/utils/working-directory.ts | 36 +++--- tools/ui/src/routes/+error.svelte | 3 +- tools/ui/src/routes/+layout.svelte | 16 ++- tools/ui/src/routes/search/+page.svelte | 2 +- .../tests/client/apikey-splash.svelte.test.ts | 2 +- .../chat-form-enter-code-block.svelte.test.ts | 2 +- ...ettings-registry-invariants.svelte.test.ts | 3 +- ...tings-render-keys-migration.svelte.test.ts | 2 +- .../client/ui-settings-sync.svelte.test.ts | 2 +- tools/ui/tests/unit/agentic-strip.test.ts | 2 +- .../tests/unit/mcp-override-fallback.test.ts | 3 +- .../ui/tests/unit/mcp-servers-default.test.ts | 4 +- tools/ui/tests/unit/mcp-service.test.ts | 8 +- .../unit/parse-mcp-server-settings.test.ts | 2 +- tools/ui/tests/unit/sanitize-headers.test.ts | 8 +- tools/ui/tests/unit/uri-template.test.ts | 8 +- tools/ui/tests/unit/working-directory.test.ts | 10 +- tools/ui/vite.config.ts | 2 +- 209 files changed, 1101 insertions(+), 1137 deletions(-) rename tools/ui/src/lib/constants/{agentic.ts => agentic.constants.ts} (77%) rename tools/ui/src/lib/constants/{api-endpoints.ts => api-endpoints.constants.ts} (100%) rename tools/ui/src/lib/constants/{app.ts => app.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/attachment-labels.ts rename tools/ui/src/lib/constants/{attachment-menu.ts => attachment-menu.constants.ts} (75%) rename tools/ui/src/lib/constants/{auto-scroll.ts => auto-scroll.constants.ts} (100%) rename tools/ui/src/lib/constants/{binary-detection.ts => binary-detection.constants.ts} (100%) rename tools/ui/src/lib/constants/{built-in-tools.ts => built-in-tools.constants.ts} (82%) create mode 100644 tools/ui/src/lib/constants/cache.constants.ts delete mode 100644 tools/ui/src/lib/constants/cache.ts rename tools/ui/src/lib/constants/{chat-form.ts => chat-form.constants.ts} (100%) rename tools/ui/src/lib/constants/{cli-flags.ts => cli-flags.constants.ts} (100%) create mode 100644 tools/ui/src/lib/constants/code-block.constants.ts delete mode 100644 tools/ui/src/lib/constants/code-blocks.ts delete mode 100644 tools/ui/src/lib/constants/code.ts create mode 100644 tools/ui/src/lib/constants/content-detection.constants.ts rename tools/ui/src/lib/constants/{context-gauge-popup.ts => context-gauge-popup.constants.ts} (100%) rename tools/ui/src/lib/constants/{context-keys.ts => context-keys.constants.ts} (100%) rename tools/ui/src/lib/constants/{control-actions.ts => control-actions.constants.ts} (73%) rename tools/ui/src/lib/constants/{conversation-import.ts => conversation-import.constants.ts} (100%) rename tools/ui/src/lib/constants/{css-classes.ts => css-classes.constants.ts} (100%) rename tools/ui/src/lib/constants/{database.ts => database.constants.ts} (93%) rename tools/ui/src/lib/constants/{diagram-blocks.ts => diagram-blocks.constants.ts} (100%) rename tools/ui/src/lib/constants/{error.ts => error.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/floating-ui-constraints.ts rename tools/ui/src/lib/constants/{formatters.ts => formatters.constants.ts} (100%) create mode 100644 tools/ui/src/lib/constants/headers.constants.ts rename tools/ui/src/lib/constants/{icons.ts => icons.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/image-size.ts create mode 100644 tools/ui/src/lib/constants/image.constants.ts delete mode 100644 tools/ui/src/lib/constants/jpeg-exif.ts rename tools/ui/src/lib/constants/{key-value-pairs.ts => key-value-pairs.constants.ts} (100%) rename tools/ui/src/lib/constants/{latex-protection.ts => latex-protection.constants.ts} (100%) rename tools/ui/src/lib/constants/{literal-html.ts => literal-html.constants.ts} (56%) create mode 100644 tools/ui/src/lib/constants/markdown.constants.ts delete mode 100644 tools/ui/src/lib/constants/markdown.ts rename tools/ui/src/lib/constants/{max-bundle-size.ts => max-bundle-size.constants.ts} (100%) rename tools/ui/src/lib/constants/{mcp-form.ts => mcp-form.constants.ts} (100%) rename tools/ui/src/lib/constants/{mcp-resource.ts => mcp-resource.constants.ts} (100%) rename tools/ui/src/lib/constants/{mcp.ts => mcp.constants.ts} (57%) rename tools/ui/src/lib/constants/{mention-badge.ts => mention-badge.constants.ts} (100%) rename tools/ui/src/lib/constants/{mermaid-blocks.ts => mermaid-blocks.constants.ts} (100%) create mode 100644 tools/ui/src/lib/constants/message-export.constants.ts delete mode 100644 tools/ui/src/lib/constants/message-export.ts create mode 100644 tools/ui/src/lib/constants/model-id.constants.ts delete mode 100644 tools/ui/src/lib/constants/model-id.ts rename tools/ui/src/lib/constants/{model-loading.ts => model-loading.constants.ts} (100%) rename tools/ui/src/lib/constants/{path-display.ts => path-display.constants.ts} (100%) rename tools/ui/src/lib/constants/{precision.ts => precision.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/processing-info.ts rename tools/ui/src/lib/constants/{pwa.ts => pwa.constants.ts} (99%) delete mode 100644 tools/ui/src/lib/constants/reasoning-effort-tokens.ts rename tools/ui/src/lib/constants/{reasoning-effort.ts => reasoning-effort.constants.ts} (71%) rename tools/ui/src/lib/constants/{recommended-mcp-servers.ts => recommended-mcp-servers.constants.ts} (100%) rename tools/ui/src/lib/constants/{routes.ts => routes.constants.ts} (100%) create mode 100644 tools/ui/src/lib/constants/sandbox.constants.ts rename tools/ui/src/lib/constants/{settings-keys.ts => settings-keys.constants.ts} (100%) rename tools/ui/src/lib/constants/{settings-registry.ts => settings-registry.constants.ts} (97%) create mode 100644 tools/ui/src/lib/constants/special-characters.constants.ts rename tools/ui/src/lib/constants/{storage.ts => storage.constants.ts} (100%) rename tools/ui/src/lib/constants/{sse.ts => stream.constants.ts} (53%) delete mode 100644 tools/ui/src/lib/constants/stream.ts rename tools/ui/src/lib/constants/{supported-file-types.ts => supported-file-types.constants.ts} (100%) create mode 100644 tools/ui/src/lib/constants/svg-blocks.constants.ts delete mode 100644 tools/ui/src/lib/constants/svg-blocks.ts rename tools/ui/src/lib/constants/{table-html-restorer.ts => table-html-restorer.constants.ts} (100%) rename tools/ui/src/lib/constants/{title-generation.ts => title-generation.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/tools.ts delete mode 100644 tools/ui/src/lib/constants/tooltip-config.ts rename tools/ui/src/lib/constants/{ui.ts => ui.constants.ts} (57%) rename tools/ui/src/lib/constants/{uri-template.ts => uri-template.constants.ts} (71%) rename tools/ui/src/lib/constants/{url.ts => url.constants.ts} (100%) delete mode 100644 tools/ui/src/lib/constants/viewport.ts create mode 100644 tools/ui/src/lib/constants/working-directory.constants.ts delete mode 100644 tools/ui/src/lib/constants/working-directory.ts create mode 100644 tools/ui/src/lib/types/navigation.d.ts create mode 100644 tools/ui/src/lib/utils/built-in-tools.ts rename tools/ui/src/lib/{constants => utils}/chat-commands.ts (65%) rename tools/ui/src/lib/{constants/sandbox.ts => utils/sandbox-tool.ts} (84%) diff --git a/tools/ui/pwa-assets-dark.config.ts b/tools/ui/pwa-assets-dark.config.ts index 1446b47a82a4..4d8114ee76df 100644 --- a/tools/ui/pwa-assets-dark.config.ts +++ b/tools/ui/pwa-assets-dark.config.ts @@ -1,5 +1,5 @@ import { writeThemeFavicons } from './scripts/favicon-colorize'; -import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa'; +import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa.constants'; import { defineConfig } from '@vite-pwa/assets-generator/config'; writeThemeFavicons(FAVICON_COLORS.LIGHT, FAVICON_COLORS.DARK, { diff --git a/tools/ui/pwa-assets.config.ts b/tools/ui/pwa-assets.config.ts index 5fed0a595743..f9f8662a20af 100644 --- a/tools/ui/pwa-assets.config.ts +++ b/tools/ui/pwa-assets.config.ts @@ -4,7 +4,7 @@ import { PWA_ASSET_GENERATOR, PWA_GENERATOR_DEVICES, THEME_COLORS -} from './src/lib/constants/pwa'; +} from './src/lib/constants/pwa.constants'; import { SplashOrientation } from './src/lib/enums/splash.enums'; import { combinePresetAndAppleSplashScreens, diff --git a/tools/ui/scripts/vite-plugin-build-info.ts b/tools/ui/scripts/vite-plugin-build-info.ts index 80023863057b..ec864e8d03f9 100644 --- a/tools/ui/scripts/vite-plugin-build-info.ts +++ b/tools/ui/scripts/vite-plugin-build-info.ts @@ -1,4 +1,4 @@ -import { BUILD_CONFIG } from '../src/lib/constants/pwa'; +import { BUILD_CONFIG } from '../src/lib/constants/pwa.constants'; import { existsSync, writeFileSync } from 'node:fs'; import { resolve } from 'path'; import type { Plugin } from 'vite'; diff --git a/tools/ui/scripts/vite-plugin-relativize-base.ts b/tools/ui/scripts/vite-plugin-relativize-base.ts index f8eac1d66555..0e47741ae94d 100644 --- a/tools/ui/scripts/vite-plugin-relativize-base.ts +++ b/tools/ui/scripts/vite-plugin-relativize-base.ts @@ -1,4 +1,4 @@ -import { BUILD_CONFIG } from '../src/lib/constants/pwa'; +import { BUILD_CONFIG } from '../src/lib/constants/pwa.constants'; import { existsSync, readFileSync, writeFileSync } from 'node:fs'; import { resolve } from 'path'; import type { Plugin } from 'vite'; diff --git a/tools/ui/scripts/vite-plugin-splash-screen.ts b/tools/ui/scripts/vite-plugin-splash-screen.ts index 45d931bab36c..62b7a063acde 100644 --- a/tools/ui/scripts/vite-plugin-splash-screen.ts +++ b/tools/ui/scripts/vite-plugin-splash-screen.ts @@ -1,5 +1,10 @@ -import { NEWLINE, TAB } from '../src/lib/constants/code'; -import { APPLE_DEVICES, BUILD_CONFIG, REGEX_PATTERNS, SPLASH_LINK } from '../src/lib/constants/pwa'; +import { + APPLE_DEVICES, + BUILD_CONFIG, + REGEX_PATTERNS, + SPLASH_LINK +} from '../src/lib/constants/pwa.constants'; +import { NEWLINE, TAB } from '../src/lib/constants/special-characters.constants'; import { SplashOrientation } from '../src/lib/enums/splash.enums'; import type { SplashDimensions } from '../src/lib/types'; import { existsSync, readdirSync, readFileSync, writeFileSync } from 'node:fs'; diff --git a/tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte b/tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte index 7655f18e804e..2d54df89d67c 100644 --- a/tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte +++ b/tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte @@ -1,7 +1,7 @@ diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte index ae5d91beee5e..cae37d6f8a8b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte @@ -1,7 +1,7 @@ diff --git a/tools/ui/src/lib/constants/agentic.ts b/tools/ui/src/lib/constants/agentic.constants.ts similarity index 77% rename from tools/ui/src/lib/constants/agentic.ts rename to tools/ui/src/lib/constants/agentic.constants.ts index 582572a29afc..e57104e8a83f 100644 --- a/tools/ui/src/lib/constants/agentic.ts +++ b/tools/ui/src/lib/constants/agentic.constants.ts @@ -5,22 +5,14 @@ export const ATTACHMENT_SAVED_REGEX = /\[Attachment saved: ([^\]]+)\]/; // JSON detection: trimmed content opens with an object or array literal. export const TOOL_RESULT_JSON_OPEN_REGEX = /^[[{]/; -// Markdown structural markers used by `looksLikeMarkdown`. Inline / line-level. -export const MARKDOWN_CODE_FENCE_REGEX = /^(```|~~~)/m; -export const MARKDOWN_ATX_HEADING_REGEX = /^#{1,6}\s+\S/; -export const MARKDOWN_BLOCKQUOTE_REGEX = /^>\s+\S/; -export const MARKDOWN_LIST_BULLET_REGEX = /^\s*[-*+]\s+\S/; -export const MARKDOWN_LIST_NUMBERED_REGEX = /^\s*\d+[.)]\s+\S/; -export const MARKDOWN_LINK_REGEX = /\[[^\]\n]+\]\([^)\s]+\)/; -export const MARKDOWN_BOLD_REGEX = /\*\*[^*\n]+\*\*|__[^_\n]+__/; -export const MARKDOWN_TABLE_SEPARATOR_REGEX = /^\s*\|?[\s:|-]+\|?\s*$/; - // Search-summary wire format used by file-glob and grep tools: // // --- // Total matches: N -export const SEARCH_SUMMARY_SEPARATOR = '---\n'; -export const SEARCH_SUMMARY_TOTAL_REGEX = /Total matches:\s*(\d+)/; +export const SEARCH_SUMMARY = { + SEPARATOR: '---\n', + TOTAL_REGEX: /Total matches:\s*(\d+)/ +} as const; // Separator rendered between stats in the tool-result footer (e.g. between a // result message and the byte/edit count). Plain ASCII spaces bracket a hyphen diff --git a/tools/ui/src/lib/constants/api-endpoints.ts b/tools/ui/src/lib/constants/api-endpoints.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/api-endpoints.ts rename to tools/ui/src/lib/constants/api-endpoints.constants.ts diff --git a/tools/ui/src/lib/constants/app.ts b/tools/ui/src/lib/constants/app.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/app.ts rename to tools/ui/src/lib/constants/app.constants.ts diff --git a/tools/ui/src/lib/constants/attachment-labels.ts b/tools/ui/src/lib/constants/attachment-labels.ts deleted file mode 100644 index be9999c0f9ee..000000000000 --- a/tools/ui/src/lib/constants/attachment-labels.ts +++ /dev/null @@ -1,4 +0,0 @@ -export const ATTACHMENT_LABEL_FILE = 'File'; -export const ATTACHMENT_LABEL_PDF_FILE = 'PDF File'; -export const ATTACHMENT_LABEL_MCP_PROMPT = 'MCP Prompt'; -export const ATTACHMENT_LABEL_MCP_RESOURCE = 'MCP Resource'; diff --git a/tools/ui/src/lib/constants/attachment-menu.ts b/tools/ui/src/lib/constants/attachment-menu.constants.ts similarity index 75% rename from tools/ui/src/lib/constants/attachment-menu.ts rename to tools/ui/src/lib/constants/attachment-menu.constants.ts index 1cd7f9ba5ae0..62e03bea6e24 100644 --- a/tools/ui/src/lib/constants/attachment-menu.ts +++ b/tools/ui/src/lib/constants/attachment-menu.constants.ts @@ -1,33 +1,12 @@ import { FolderOpen, MessageSquare, Zap } from '@lucide/svelte'; -import { FILE_TYPE_ICONS } from '$lib/constants/icons'; +import { FILE_TYPE_ICONS } from '$lib/constants'; import { AttachmentAction, AttachmentItemEnabledWhen, AttachmentItemVisibleWhen, AttachmentMenuItemId } from '$lib/enums'; -import type { Component } from 'svelte'; - -export interface AttachmentMenuItem { - /** Unique identifier for the item */ - id: AttachmentMenuItemId; - /** Display label */ - label: string; - /** Lucide icon component */ - icon: Component; - /** Extra CSS class applied to the item (e.g. for test selectors) */ - class?: string; - /** Whether the item requires a specific modality to be enabled */ - enabledWhen?: AttachmentItemEnabledWhen; - /** Tooltip shown when the item is disabled */ - disabledTooltip?: string; - /** Callback key on the Props interface to invoke when clicked */ - action: AttachmentAction; - /** Whether the item is only shown when a specific capability is present */ - visibleWhen?: AttachmentItemVisibleWhen; - /** Whether this item has a tooltip even when enabled (uses dynamic text) */ - hasEnabledTooltip?: boolean; -} +import type { AttachmentMenuItem } from '$lib/types'; /** * File attachment menu items shown in both the desktop dropdown and mobile sheet. diff --git a/tools/ui/src/lib/constants/auto-scroll.ts b/tools/ui/src/lib/constants/auto-scroll.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/auto-scroll.ts rename to tools/ui/src/lib/constants/auto-scroll.constants.ts diff --git a/tools/ui/src/lib/constants/binary-detection.ts b/tools/ui/src/lib/constants/binary-detection.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/binary-detection.ts rename to tools/ui/src/lib/constants/binary-detection.constants.ts diff --git a/tools/ui/src/lib/constants/built-in-tools.ts b/tools/ui/src/lib/constants/built-in-tools.constants.ts similarity index 82% rename from tools/ui/src/lib/constants/built-in-tools.ts rename to tools/ui/src/lib/constants/built-in-tools.constants.ts index 5bd24ffe8397..679c61459232 100644 --- a/tools/ui/src/lib/constants/built-in-tools.ts +++ b/tools/ui/src/lib/constants/built-in-tools.constants.ts @@ -20,13 +20,7 @@ import { Terminal } from '@lucide/svelte'; import { BuiltInTool, ToolSource } from '$lib/enums'; -import type { Component } from 'svelte'; - -export interface BuiltinToolUiEntry { - icon: Component; - label: string; - source: ToolSource.BUILTIN | ToolSource.FRONTEND; -} +import type { BuiltinToolUiEntry } from '$lib/types'; export const BUILTIN_TOOL_UI: Readonly> = { [BuiltInTool.EDIT_FILE]: { icon: FilePen, label: 'Edit file', source: ToolSource.BUILTIN }, @@ -56,9 +50,3 @@ export const BUILTIN_TOOL_UI: Readonly> }, [BuiltInTool.WRITE_FILE]: { icon: FilePlus, label: 'Write file', source: ToolSource.BUILTIN } } as const; - -export function getBuiltinToolUi(toolName: string | undefined): BuiltinToolUiEntry | null { - if (!toolName) return null; - - return (BUILTIN_TOOL_UI as Record)[toolName] ?? null; -} diff --git a/tools/ui/src/lib/constants/cache.constants.ts b/tools/ui/src/lib/constants/cache.constants.ts new file mode 100644 index 000000000000..b60792d99524 --- /dev/null +++ b/tools/ui/src/lib/constants/cache.constants.ts @@ -0,0 +1,44 @@ +/** + * Cache configuration constants + */ + +/** + * Default cache limits when no per-cache overrides are given. + */ +export const CACHE = { + /** Default maximum number of entries in a cache */ + DEFAULT_MAX_ENTRIES: 100, + /** Default TTL (Time-To-Live) for cache entries in milliseconds (5 minutes) */ + DEFAULT_TTL_MS: 5 * 60 * 1000 +} as const; + +/** + * TTL and size for the model props cache. + * Props don't change frequently, so we can cache them longer. + */ +export const MODEL_PROPS_CACHE = { + /** Maximum number of model props to cache */ + MAX_ENTRIES: 50, + /** TTL for model props cache entries in milliseconds (10 minutes) */ + TTL_MS: 10 * 60 * 1000 +} as const; + +/** + * TTL and size for the MCP resource cache. + */ +export const MCP_RESOURCE_CACHE = { + /** Maximum number of MCP resources to cache */ + MAX_ENTRIES: 50, + /** TTL for MCP resource cache entries in milliseconds (5 minutes) */ + TTL_MS: 5 * 60 * 1000 +} as const; + +/** + * Limits for pruning inactive conversation states held in memory. + */ +export const INACTIVE_CONVERSATION = { + /** Maximum age (in ms) for inactive conversation states before cleanup (30 minutes) */ + MAX_AGE_MS: 30 * 60 * 1000, + /** Maximum number of inactive conversation states to keep in memory */ + MAX_STATES: 10 +} as const; diff --git a/tools/ui/src/lib/constants/cache.ts b/tools/ui/src/lib/constants/cache.ts deleted file mode 100644 index 07fe8683414a..000000000000 --- a/tools/ui/src/lib/constants/cache.ts +++ /dev/null @@ -1,54 +0,0 @@ -/** - * Cache configuration constants - */ - -/** - * Default TTL (Time-To-Live) for cache entries in milliseconds - * @default 5 minutes - */ -export const DEFAULT_CACHE_TTL_MS = 5 * 60 * 1000; - -/** - * Default maximum number of entries in a cache - * @default 100 - */ -export const DEFAULT_CACHE_MAX_ENTRIES = 100; - -/** - * TTL for model props cache in milliseconds - * Props don't change frequently, so we can cache them longer - * @default 10 minutes - */ -export const MODEL_PROPS_CACHE_TTL_MS = 10 * 60 * 1000; - -/** - * Maximum number of model props to cache - * @default 50 - */ -export const MODEL_PROPS_CACHE_MAX_ENTRIES = 50; - -/** - * Maximum number of MCP resources to cache - * @default 50 - */ -export const MCP_RESOURCE_CACHE_MAX_ENTRIES = 50; - -/** - * TTL for MCP resource cache entries in milliseconds - * @default 5 minutes - */ -export const MCP_RESOURCE_CACHE_TTL_MS = 5 * 60 * 1000; - -/** - * Maximum number of inactive conversation states to keep in memory - * States for conversations beyond this limit will be cleaned up - * @default 10 - */ -export const MAX_INACTIVE_CONVERSATION_STATES = 10; - -/** - * Maximum age (in ms) for inactive conversation states before cleanup - * States older than this will be removed during cleanup - * @default 30 minutes - */ -export const INACTIVE_CONVERSATION_STATE_MAX_AGE_MS = 30 * 60 * 1000; diff --git a/tools/ui/src/lib/constants/chat-form.ts b/tools/ui/src/lib/constants/chat-form.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/chat-form.ts rename to tools/ui/src/lib/constants/chat-form.constants.ts diff --git a/tools/ui/src/lib/constants/cli-flags.ts b/tools/ui/src/lib/constants/cli-flags.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/cli-flags.ts rename to tools/ui/src/lib/constants/cli-flags.constants.ts diff --git a/tools/ui/src/lib/constants/code-block.constants.ts b/tools/ui/src/lib/constants/code-block.constants.ts new file mode 100644 index 000000000000..05db575f9ea5 --- /dev/null +++ b/tools/ui/src/lib/constants/code-block.constants.ts @@ -0,0 +1,50 @@ +// Constants for the markdown code-block renderer: language/fence handling and CSS classes. + +/** Parsing and escaping helpers for the markdown code-block renderer. */ +export const CODE_BLOCK = { + AMPERSAND_REGEX: /&/g, + /** Language fallback used when no language is specified. */ + DEFAULT_LANGUAGE: 'text', + /** Matches opening/closing markdown code fences. */ + FENCE_PATTERN: /^```|\n```/g, + GT_REGEX: />/g, + /** Matches the language specifier at the start of a code fence. */ + LANG_PATTERN: /^(\w*)\n?/, + LT_REGEX: //g; -export const FENCE_PATTERN = /^```|\n```/g; - -// Whitespace-only empty lines (between start of string and first non-empty line). -// Used by trimCodePadding to drop leading/trailing phantom blank rows from LLM -// payload wrappers without touching internal blank lines. -export const TRIM_LEADING_PADDING_REGEX = /^(?:[ \t]*\n)+/; -export const TRIM_TRAILING_PADDING_REGEX = /(?:\n[ \t]*)+$/; - -// Matches either Unix or Windows path separators so `String.split(REGEX)` can -// recover the trailing file-name segment from either `/foo/bar.txt` or -// `C:\foo\bar.txt`. Used wherever a parameter accepts a user-supplied path. -export const FILE_PATH_SEPARATOR_REGEX = /[\\/]/; - -// Separates a file name from its extension, e.g. the '.' in `cover.png`. -export const FILE_EXTENSION_SEPARATOR = '.'; - -// Matches the `text:` prefix that file-type identifiers use to denote a -// plain-text language (e.g. `text:typescript`). Used by tool-call renderers -// to recover the underlying highlight.js language. -export const TEXT_LANGUAGE_PREFIX_REGEX = /^text:/; diff --git a/tools/ui/src/lib/constants/content-detection.constants.ts b/tools/ui/src/lib/constants/content-detection.constants.ts new file mode 100644 index 000000000000..c5c05819a503 --- /dev/null +++ b/tools/ui/src/lib/constants/content-detection.constants.ts @@ -0,0 +1,20 @@ +/** + * String patterns for detecting content kind from MIME types and URIs. + * Used with startsWith/includes checks, not as discriminated values. + */ + +export const MIME_TYPE_PREFIXES = { + IMAGE: 'image/', + TEXT: 'text' +} as const; + +export const MIME_TYPE_SUBSTRINGS = { + JAVASCRIPT: 'javascript', + JSON: 'json', + TYPESCRIPT: 'typescript' +} as const; + +export const URI_PATTERNS = { + DATABASE_KEYWORD: 'database', + DATABASE_SCHEME: 'db://' +} as const; diff --git a/tools/ui/src/lib/constants/context-gauge-popup.ts b/tools/ui/src/lib/constants/context-gauge-popup.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/context-gauge-popup.ts rename to tools/ui/src/lib/constants/context-gauge-popup.constants.ts diff --git a/tools/ui/src/lib/constants/context-keys.ts b/tools/ui/src/lib/constants/context-keys.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/context-keys.ts rename to tools/ui/src/lib/constants/context-keys.constants.ts diff --git a/tools/ui/src/lib/constants/control-actions.ts b/tools/ui/src/lib/constants/control-actions.constants.ts similarity index 73% rename from tools/ui/src/lib/constants/control-actions.ts rename to tools/ui/src/lib/constants/control-actions.constants.ts index 935ae9542a36..c8ebf701b1b7 100644 --- a/tools/ui/src/lib/constants/control-actions.ts +++ b/tools/ui/src/lib/constants/control-actions.constants.ts @@ -3,5 +3,3 @@ export const CONTROL_ACTION = { END_REASONING: 'reasoning_end' } as const; - -export type ControlAction = (typeof CONTROL_ACTION)[keyof typeof CONTROL_ACTION]; diff --git a/tools/ui/src/lib/constants/conversation-import.ts b/tools/ui/src/lib/constants/conversation-import.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/conversation-import.ts rename to tools/ui/src/lib/constants/conversation-import.constants.ts diff --git a/tools/ui/src/lib/constants/css-classes.ts b/tools/ui/src/lib/constants/css-classes.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/css-classes.ts rename to tools/ui/src/lib/constants/css-classes.constants.ts diff --git a/tools/ui/src/lib/constants/database.ts b/tools/ui/src/lib/constants/database.constants.ts similarity index 93% rename from tools/ui/src/lib/constants/database.ts rename to tools/ui/src/lib/constants/database.constants.ts index 95e698f40012..f2c96103932d 100644 --- a/tools/ui/src/lib/constants/database.ts +++ b/tools/ui/src/lib/constants/database.constants.ts @@ -5,7 +5,7 @@ * naming changes. */ -import { STORAGE_APP_NAME } from './storage'; +import { STORAGE_APP_NAME } from './storage.constants'; /** IndexedDB database name */ export const DB_NAME = STORAGE_APP_NAME; diff --git a/tools/ui/src/lib/constants/diagram-blocks.ts b/tools/ui/src/lib/constants/diagram-blocks.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/diagram-blocks.ts rename to tools/ui/src/lib/constants/diagram-blocks.constants.ts diff --git a/tools/ui/src/lib/constants/error.ts b/tools/ui/src/lib/constants/error.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/error.ts rename to tools/ui/src/lib/constants/error.constants.ts diff --git a/tools/ui/src/lib/constants/floating-ui-constraints.ts b/tools/ui/src/lib/constants/floating-ui-constraints.ts deleted file mode 100644 index 003fc77acb08..000000000000 --- a/tools/ui/src/lib/constants/floating-ui-constraints.ts +++ /dev/null @@ -1,2 +0,0 @@ -export const VIEWPORT_GUTTER = 8; -export const MENU_OFFSET = 6; diff --git a/tools/ui/src/lib/constants/formatters.ts b/tools/ui/src/lib/constants/formatters.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/formatters.ts rename to tools/ui/src/lib/constants/formatters.constants.ts diff --git a/tools/ui/src/lib/constants/headers.constants.ts b/tools/ui/src/lib/constants/headers.constants.ts new file mode 100644 index 000000000000..f40df05710f3 --- /dev/null +++ b/tools/ui/src/lib/constants/headers.constants.ts @@ -0,0 +1,35 @@ +/** HTTP header handling for API and MCP requests. */ +export const HEADERS = { + /** Canonical casing for the Authorization header (RFC 7235) */ + AUTHORIZATION: 'Authorization', + /** Bearer scheme prefix used for Authorization headers (RFC 6750) */ + BEARER: 'Bearer ', + /** Content-Type HTTP header name */ + CONTENT_TYPE: 'Content-Type', + /** Partial-redaction rules for MCP headers: header name -> visible trailing chars */ + PARTIAL_REDACT: new Map([['mcp-session-id', 5]]), + + /** Header names whose values should be redacted in diagnostic logs */ + REDACTED: new Set([ + 'authorization', + 'api-key', + 'cookie', + 'mcp-session-id', + 'proxy-authorization', + 'set-cookie', + 'x-auth-token', + 'x-api-key' + ]), + + /** Header carrying the stream-session identity (conversation id, optionally with a model suffix) */ + X_CONVERSATION_ID_HEADER: 'X-Conversation-Id', + + /** Header asking the server to encode a tool's output differently, e.g. read_file returning base64. */ + X_RESP_TYPE_HEADER: 'x-resp-type', + + /** Header carrying the working directory a tool call runs in; the model cannot override it */ + X_TOOL_CWD_HEADER: 'x-tool-cwd' +}; + +/** `X_RESP_TYPE_HEADER` value that makes read_file return raw bytes as base64 instead of text. */ +export const RESP_TYPE_BASE64 = 'base64'; diff --git a/tools/ui/src/lib/constants/icons.ts b/tools/ui/src/lib/constants/icons.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/icons.ts rename to tools/ui/src/lib/constants/icons.constants.ts diff --git a/tools/ui/src/lib/constants/image-size.ts b/tools/ui/src/lib/constants/image-size.ts deleted file mode 100644 index 8a7f921fa019..000000000000 --- a/tools/ui/src/lib/constants/image-size.ts +++ /dev/null @@ -1,3 +0,0 @@ -export const MEGAPIXELS_TO_PIXELS = 1_000_000; - -export const HEIC_JPEG_QUALITY = 0.85; diff --git a/tools/ui/src/lib/constants/image.constants.ts b/tools/ui/src/lib/constants/image.constants.ts new file mode 100644 index 000000000000..53a90eaa4e19 --- /dev/null +++ b/tools/ui/src/lib/constants/image.constants.ts @@ -0,0 +1,32 @@ +/** Image handling constants */ + +export const IMAGE = { + /** JPEG quality used when transcoding HEIC images. */ + HEIC_JPEG_QUALITY: 0.85, + /** Unit conversion: pixels per megapixel. */ + MEGAPIXELS_TO_PIXELS: 1_000_000 +} as const; + +/** + * JPEG and EXIF binary format constants for orientation parsing. + */ +export const EXIF = { + /** APP1 segment marker byte, carries the EXIF payload */ + APP1_MARKER: 0xe1, + /** "Exif" signature opening the APP1 payload, big endian uint32 */ + EXIF_SIGNATURE: 0x45786966, + /** Size in bytes of one IFD directory entry */ + IFD_ENTRY_SIZE: 12, + /** JPEG start of image marker */ + JPEG_SOI_MARKER: 0xffd8, + /** EXIF tag id holding the orientation value */ + ORIENTATION_TAG: 0x0112, + /** Bytes of file prefix to scan, the APP1 EXIF segment sits near the start */ + SCAN_BYTE_LIMIT: 128 * 1024, + /** Start of scan marker byte, compressed data begins and no EXIF follows */ + SOS_MARKER: 0xda, + /** TIFF byte order mark for little endian ("II") */ + TIFF_LITTLE_ENDIAN: 0x4949, + /** TIFF magic number following the byte order mark */ + TIFF_MAGIC: 42 +} as const; diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index 357a33a6241e..17e5f4d00ab5 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -1,67 +1,61 @@ // Central constants export file // All constants should be imported from '$lib/constants' -export * from './agentic'; -export * from './api-endpoints'; -export * from './app'; -export * from './attachment-labels'; -export * from './database'; -export * from './reasoning-effort'; -export * from './reasoning-effort-tokens'; -export * from './recommended-mcp-servers'; -export * from './storage'; -export * from './attachment-menu'; -export * from './auto-scroll'; -export * from './context-gauge-popup'; -export * from './conversation-import'; -export * from './binary-detection'; -export * from './built-in-tools'; -export * from './cache'; -export * from './chat-form'; -export * from './chat-commands'; -export * from './cli-flags'; -export * from './code-blocks'; -export * from './icons'; -export * from './code'; -export * from './context-keys'; -export * from './control-actions'; -export * from './css-classes'; -export * from './floating-ui-constraints'; -export * from './formatters'; -export * from './key-value-pairs'; -export * from './icons'; -export * from './latex-protection'; -export * from './literal-html'; -export * from './markdown'; -export * from './mermaid-blocks'; -export * from './svg-blocks'; -export * from './diagram-blocks'; -export * from './max-bundle-size'; -export * from './mcp'; -export * from './mcp-form'; -export * from './mcp-resource'; -export * from './mention-badge'; -export * from './message-export'; -export * from './path-display'; -export * from './model-id'; -export * from './model-loading'; -export * from './sse'; -export * from './precision'; -export * from './processing-info'; -export * from './pwa'; +export * from './agentic.constants'; +export * from './api-endpoints.constants'; +export * from './app.constants'; +export * from './database.constants'; +export * from './reasoning-effort.constants'; +export * from './recommended-mcp-servers.constants'; +export * from './storage.constants'; +export * from './icons.constants'; +export * from './attachment-menu.constants'; +export * from './auto-scroll.constants'; +export * from './context-gauge-popup.constants'; +export * from './conversation-import.constants'; +export * from './binary-detection.constants'; +export * from './content-detection.constants'; +export * from './built-in-tools.constants'; +export * from './cache.constants'; +export * from './chat-form.constants'; +export * from './cli-flags.constants'; +export * from './code-block.constants'; +export * from './context-keys.constants'; +export * from './control-actions.constants'; +export * from './css-classes.constants'; +export * from './formatters.constants'; +export * from './headers.constants'; +export * from './key-value-pairs.constants'; +export * from './latex-protection.constants'; +export * from './literal-html.constants'; +export * from './markdown.constants'; +export * from './mermaid-blocks.constants'; +export * from './svg-blocks.constants'; +export * from './diagram-blocks.constants'; +export * from './max-bundle-size.constants'; +export * from './error.constants'; +export * from './image.constants'; +export * from './mcp.constants'; +export * from './mcp-form.constants'; +export * from './mcp-resource.constants'; +export * from './mention-badge.constants'; +export * from './message-export.constants'; +export * from './path-display.constants'; +export * from './model-id.constants'; +export * from './model-loading.constants'; +export * from './precision.constants'; +export * from './pwa.constants'; +export * from './routes.constants'; +export * from './sandbox.constants'; +export * from './settings-keys.constants'; +export * from './settings-registry.constants'; +export * from './special-characters.constants'; +export * from './stream.constants'; +export * from './supported-file-types.constants'; +export * from './table-html-restorer.constants'; +export * from './title-generation.constants'; +export * from './ui.constants'; +export * from './uri-template.constants'; +export * from './url.constants'; +export * from './working-directory.constants'; export * from './read-media'; -export * from './routes'; -export * from './sandbox'; -export * from './settings-keys'; -export * from './settings-registry'; -export * from './stream'; -export * from './supported-file-types'; -export * from './table-html-restorer'; -export * from './title-generation'; -export * from './tools'; -export * from './tooltip-config'; -export * from './ui'; -export * from './uri-template'; -export * from './url'; -export * from './viewport'; -export * from './working-directory'; diff --git a/tools/ui/src/lib/constants/jpeg-exif.ts b/tools/ui/src/lib/constants/jpeg-exif.ts deleted file mode 100644 index 5b2591b04b18..000000000000 --- a/tools/ui/src/lib/constants/jpeg-exif.ts +++ /dev/null @@ -1,30 +0,0 @@ -/** - * JPEG and EXIF binary format constants for orientation parsing. - */ - -/** Bytes of file prefix to scan, the APP1 EXIF segment sits near the start */ -export const EXIF_SCAN_BYTE_LIMIT = 128 * 1024; - -/** JPEG start of image marker */ -export const JPEG_SOI_MARKER = 0xffd8; - -/** APP1 segment marker byte, carries the EXIF payload */ -export const APP1_MARKER = 0xe1; - -/** Start of scan marker byte, compressed data begins and no EXIF follows */ -export const SOS_MARKER = 0xda; - -/** "Exif" signature opening the APP1 payload, big endian uint32 */ -export const EXIF_SIGNATURE = 0x45786966; - -/** TIFF byte order mark for little endian ("II") */ -export const TIFF_LITTLE_ENDIAN = 0x4949; - -/** TIFF magic number following the byte order mark */ -export const TIFF_MAGIC = 42; - -/** EXIF tag id holding the orientation value */ -export const EXIF_ORIENTATION_TAG = 0x0112; - -/** Size in bytes of one IFD directory entry */ -export const IFD_ENTRY_SIZE = 12; diff --git a/tools/ui/src/lib/constants/key-value-pairs.ts b/tools/ui/src/lib/constants/key-value-pairs.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/key-value-pairs.ts rename to tools/ui/src/lib/constants/key-value-pairs.constants.ts diff --git a/tools/ui/src/lib/constants/latex-protection.ts b/tools/ui/src/lib/constants/latex-protection.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/latex-protection.ts rename to tools/ui/src/lib/constants/latex-protection.constants.ts diff --git a/tools/ui/src/lib/constants/literal-html.ts b/tools/ui/src/lib/constants/literal-html.constants.ts similarity index 56% rename from tools/ui/src/lib/constants/literal-html.ts rename to tools/ui/src/lib/constants/literal-html.constants.ts index ed1b0cf0d90d..8efa6b5747e1 100644 --- a/tools/ui/src/lib/constants/literal-html.ts +++ b/tools/ui/src/lib/constants/literal-html.constants.ts @@ -1,5 +1,3 @@ -export const LINE_BREAK = /\r?\n/; - export const PHRASE_PARENTS = new Set([ 'paragraph', 'heading', @@ -10,6 +8,3 @@ export const PHRASE_PARENTS = new Set([ 'linkReference', 'tableCell' ]); - -export const NBSP = '\u00a0'; -export const TAB_AS_SPACES = NBSP.repeat(4); diff --git a/tools/ui/src/lib/constants/markdown.constants.ts b/tools/ui/src/lib/constants/markdown.constants.ts new file mode 100644 index 000000000000..c8289c39a1de --- /dev/null +++ b/tools/ui/src/lib/constants/markdown.constants.ts @@ -0,0 +1,17 @@ +export const IMAGE_NOT_ERROR_BOUND_SELECTOR = 'img:not([data-error-bound])'; +export const DATA_ERROR_BOUND_ATTR = 'errorBound'; +export const DATA_ERROR_HANDLED_ATTR = 'errorHandled'; +export const BOOL_TRUE_STRING = 'true'; +export const BOOL_FALSE_STRING = 'false'; + +/** Markdown structural markers used by `looksLikeMarkdown`. Inline / line-level. */ +export const MARKDOWN = { + ATX_HEADING_REGEX: /^#{1,6}\s+\S/, + BLOCKQUOTE_REGEX: /^>\s+\S/, + BOLD_REGEX: /\*\*[^*\n]+\*\*|__[^_\n]+__/, + CODE_FENCE_REGEX: /^(```|~~~)/m, + LINK_REGEX: /\[[^\]\n]+\]\([^)\s]+\)/, + LIST_BULLET_REGEX: /^\s*[-*+]\s+\S/, + LIST_NUMBERED_REGEX: /^\s*\d+[.)]\s+\S/, + TABLE_SEPARATOR_REGEX: /^\s*\|?[\s:|-]+\|?\s*$/ +} as const; diff --git a/tools/ui/src/lib/constants/markdown.ts b/tools/ui/src/lib/constants/markdown.ts deleted file mode 100644 index 1cace78a30eb..000000000000 --- a/tools/ui/src/lib/constants/markdown.ts +++ /dev/null @@ -1,5 +0,0 @@ -export const IMAGE_NOT_ERROR_BOUND_SELECTOR = 'img:not([data-error-bound])'; -export const DATA_ERROR_BOUND_ATTR = 'errorBound'; -export const DATA_ERROR_HANDLED_ATTR = 'errorHandled'; -export const BOOL_TRUE_STRING = 'true'; -export const BOOL_FALSE_STRING = 'false'; diff --git a/tools/ui/src/lib/constants/max-bundle-size.ts b/tools/ui/src/lib/constants/max-bundle-size.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/max-bundle-size.ts rename to tools/ui/src/lib/constants/max-bundle-size.constants.ts diff --git a/tools/ui/src/lib/constants/mcp-form.ts b/tools/ui/src/lib/constants/mcp-form.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/mcp-form.ts rename to tools/ui/src/lib/constants/mcp-form.constants.ts diff --git a/tools/ui/src/lib/constants/mcp-resource.ts b/tools/ui/src/lib/constants/mcp-resource.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/mcp-resource.ts rename to tools/ui/src/lib/constants/mcp-resource.constants.ts diff --git a/tools/ui/src/lib/constants/mcp.ts b/tools/ui/src/lib/constants/mcp.constants.ts similarity index 57% rename from tools/ui/src/lib/constants/mcp.ts rename to tools/ui/src/lib/constants/mcp.constants.ts index 854cbf9f76ae..11013d2cb15e 100644 --- a/tools/ui/src/lib/constants/mcp.ts +++ b/tools/ui/src/lib/constants/mcp.constants.ts @@ -36,11 +36,14 @@ export const DEFAULT_MCP_CONFIG = { export const MCP_SERVER_ID_PREFIX = 'LlamaUI-MCP-Server'; -export const MCP_RECONNECT_INITIAL_DELAY = 1000; -export const MCP_RECONNECT_BACKOFF_MULTIPLIER = 2; -export const MCP_RECONNECT_MAX_DELAY = 30000; -/** Per-attempt timeout for a single reconnection attempt before giving up and backing off. */ -export const MCP_RECONNECT_ATTEMPT_TIMEOUT_MS = 15_000; +/** Backoff policy for reconnecting to a dropped MCP server. */ +export const MCP_RECONNECT = { + /** Per-attempt timeout for a single reconnection attempt before giving up and backing off. */ + ATTEMPT_TIMEOUT_MS: 15_000, + BACKOFF_MULTIPLIER: 2, + INITIAL_DELAY: 1000, + MAX_DELAY: 30000 +}; /** Maximum number of MCP server avatars to display in the chat form */ export const MAX_DISPLAYED_MCP_AVATARS = 4; @@ -48,40 +51,20 @@ export const MAX_DISPLAYED_MCP_AVATARS = 4; /** Expected count when two theme-less icons represent a light/dark pair */ export const EXPECTED_THEMED_ICON_PAIR_COUNT = 2; -/** CORS proxy URL query parameter name */ -export const CORS_PROXY_URL_PARAM = 'url'; - -/** Header prefix for headers that should be forwarded by the CORS proxy */ -export const CORS_PROXY_HEADER_PREFIX = 'x-llama-server-proxy-header-'; - -/** Number of trailing characters to keep visible when partially redacting mcp-session-id */ -export const MCP_SESSION_ID_VISIBLE_CHARS = 5; - -/** Partial-redaction rules for MCP headers: header name -> visible trailing chars */ -export const MCP_PARTIAL_REDACT_HEADERS = new Map([ - ['mcp-session-id', MCP_SESSION_ID_VISIBLE_CHARS] -]); - -/** Bearer scheme prefix used for Authorization headers (RFC 6750) */ -export const BEARER_PREFIX = 'Bearer '; - -/** Canonical casing for the Authorization header (RFC 7235) */ -export const AUTHORIZATION_HEADER = 'Authorization'; - -/** Content-Type HTTP header name */ -export const CONTENT_TYPE_HEADER = 'Content-Type'; +/** CORS proxy connection settings */ +export const CORS_PROXY = { + /** Header prefix for headers that should be forwarded by the CORS proxy */ + HEADER_PREFIX: 'x-llama-server-proxy-header-', + /** CORS proxy URL query parameter name */ + URL_PARAM: 'url' +} as const; -/** Header names whose values should be redacted in diagnostic logs */ -export const REDACTED_HEADERS = new Set([ - 'authorization', - 'api-key', - 'cookie', - 'mcp-session-id', - 'proxy-authorization', - 'set-cookie', - 'x-auth-token', - 'x-api-key' -]); +/** Standard SSE endpoint path indicators */ +export const MCP_SSE = { + ENDPOINT: '/sse', + ENDPOINT_QUERY: '/sse?', + ENDPOINT_SLASH: '/sse/' +} as const; /** Human-readable labels for MCP transport types */ export const MCP_TRANSPORT_LABELS: Record = { @@ -96,8 +79,3 @@ export const MCP_TRANSPORT_ICONS: Record = { [MCPTransportType.STREAMABLE_HTTP]: Globe, [MCPTransportType.WEBSOCKET]: Zap }; - -/** Standard SSE endpoint path indicators */ -export const MCP_SSE_ENDPOINT = '/sse'; -export const MCP_SSE_ENDPOINT_SLASH = '/sse/'; -export const MCP_SSE_ENDPOINT_QUERY = '/sse?'; diff --git a/tools/ui/src/lib/constants/mention-badge.ts b/tools/ui/src/lib/constants/mention-badge.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/mention-badge.ts rename to tools/ui/src/lib/constants/mention-badge.constants.ts diff --git a/tools/ui/src/lib/constants/mermaid-blocks.ts b/tools/ui/src/lib/constants/mermaid-blocks.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/mermaid-blocks.ts rename to tools/ui/src/lib/constants/mermaid-blocks.constants.ts diff --git a/tools/ui/src/lib/constants/message-export.constants.ts b/tools/ui/src/lib/constants/message-export.constants.ts new file mode 100644 index 000000000000..f6c576d792b6 --- /dev/null +++ b/tools/ui/src/lib/constants/message-export.constants.ts @@ -0,0 +1,24 @@ +// Conversation exporter / filename constants + +export const EXPORT_CONV = { + // Producer marker carried by the session record of a JSONL export + HARNESS: 'llama.app', + // Length of the trimmed conversation ID in the filename + ID_TRIM_LENGTH: 8, + // Replacements to the ISO date for use in the export filename + ISO_DATE_TIME_SEPARATOR: 'T', + + ISO_DATE_TIME_SEPARATOR_REPLACEMENT: '_', + + ISO_TIME_SEPARATOR: ':', + ISO_TIME_SEPARATOR_REPLACEMENT: '-', + // Characters to keep in the ISO timestamp. 19 keeps 2026-01-01T00:00:00 + ISO_TIMESTAMP_SLICE: 19, + + MULTIPLE_UNDERSCORE_REGEX: /_+/g, + // Maximum length of the sanitized conversation name snippet + NAME_SUFFIX_MAX_LENGTH: 20, + // Replacements for making the conversation title filename-friendly + NON_ALPHANUMERIC_REGEX: /[^a-z0-9]/gi, + NONALNUM_REPLACEMENT: '_' +} as const; diff --git a/tools/ui/src/lib/constants/message-export.ts b/tools/ui/src/lib/constants/message-export.ts deleted file mode 100644 index fc4dbe259c28..000000000000 --- a/tools/ui/src/lib/constants/message-export.ts +++ /dev/null @@ -1,23 +0,0 @@ -// Conversation filename constants - -// Length of the trimmed conversation ID in the filename -export const EXPORT_CONV_ID_TRIM_LENGTH = 8; -// Maximum length of the sanitized conversation name snippet -export const EXPORT_CONV_NAME_SUFFIX_MAX_LENGTH = 20; -// Characters to keep in the ISO timestamp. 19 keeps 2026-01-01T00:00:00 -export const ISO_TIMESTAMP_SLICE_LENGTH = 19; - -// Producer marker carried by the session record of a JSONL export -export const SESSION_HARNESS = 'llama.app'; - -// Replacements for making the conversation title filename-friendly -export const NON_ALPHANUMERIC_REGEX = /[^a-z0-9]/gi; -export const EXPORT_CONV_NONALNUM_REPLACEMENT = '_'; -export const MULTIPLE_UNDERSCORE_REGEX = /_+/g; - -// Replacements to the ISO date for use in the export filename -export const ISO_DATE_TIME_SEPARATOR = 'T'; -export const ISO_DATE_TIME_SEPARATOR_REPLACEMENT = '_'; - -export const ISO_TIME_SEPARATOR = ':'; -export const ISO_TIME_SEPARATOR_REPLACEMENT = '-'; diff --git a/tools/ui/src/lib/constants/model-id.constants.ts b/tools/ui/src/lib/constants/model-id.constants.ts new file mode 100644 index 000000000000..081a13e0e69c --- /dev/null +++ b/tools/ui/src/lib/constants/model-id.constants.ts @@ -0,0 +1,43 @@ +/** + * Parsing of `org/ModelName[-tag][:quant]` style model IDs. + */ + +export const MODEL_ID = { + /** + * Matches an activated-parameter-count segment, e.g. `A10B`, `a2.4b`. + * The leading `A`/`a` distinguishes it from a regular params segment. + */ + ACTIVATED_PARAMS_RE: /^[Aa]\d+(\.\d+)?[BbMmKkTt]$/, + + /** Matches prefix for custom quantization types, e.g. `UD-Q8_K_XL`. */ + CUSTOM_QUANTIZATION_PREFIX_RE: /^UD$/i, + /** Container format segments to exclude from tags (every model uses these). */ + IGNORED_SEGMENTS: new Set(['GGUF', 'GGML']), + /** Sentinel value returned by `indexOf` when a substring is not found. */ + NOT_FOUND: -1, + + /** Separates `` from `` in a model ID, e.g. `org/ModelName`. */ + ORG_SEPARATOR: '/', + + /** + * Matches a parameter-count segment, e.g. `7B`, `1.5b`, `120M`. + * The optional leading `E` covers effective-parameter sizes, e.g. Gemma's + * `E2B`/`E4B` (MatFormer models sized by resident params). + */ + PARAMS_RE: /^[Ee]?\d+(\.\d+)?[BbMmKkTt]$/, + + /** + * Matches a quantization/precision segment, e.g. `Q4_K_M`, `IQ4_XS`, `F16`, `BF16`, `MXFP4`. + * Case-insensitive to handle both uppercase and lowercase inputs. + */ + QUANTIZATION_SEGMENT_RE: /^(I?Q\d+(_[A-Z0-9]+)*|F\d+|BF\d+|MXFP\d+(_[A-Z0-9]+)*)$/i, + + /** Separates the model path from the quantization tag, e.g. `model:Q4_K_M`. */ + QUANTIZATION_SEPARATOR: ':', + + /** Separates named segments within the model path, e.g. `ModelName-7B-GGUF`. */ + SEGMENT_SEPARATOR: '-', + + /** Matches a trailing weight file extension, e.g. `model.gguf` -> `model`. */ + WEIGHT_EXTENSION_RE: /\.(gguf|ggml)$/i +}; diff --git a/tools/ui/src/lib/constants/model-id.ts b/tools/ui/src/lib/constants/model-id.ts deleted file mode 100644 index 4108a2132126..000000000000 --- a/tools/ui/src/lib/constants/model-id.ts +++ /dev/null @@ -1,46 +0,0 @@ -/** Sentinel value returned by `indexOf` when a substring is not found. */ -export const MODEL_ID_NOT_FOUND = -1; - -/** Separates `` from `` in a model ID, e.g. `org/ModelName`. */ -export const MODEL_ID_ORG_SEPARATOR = '/'; - -/** Separates named segments within the model path, e.g. `ModelName-7B-GGUF`. */ -export const MODEL_ID_SEGMENT_SEPARATOR = '-'; - -/** Separates the model path from the quantization tag, e.g. `model:Q4_K_M`. */ -export const MODEL_ID_QUANTIZATION_SEPARATOR = ':'; - -/** - * Matches a quantization/precision segment, e.g. `Q4_K_M`, `IQ4_XS`, `F16`, `BF16`, `MXFP4`. - * Case-insensitive to handle both uppercase and lowercase inputs. - */ -export const MODEL_QUANTIZATION_SEGMENT_RE = - /^(I?Q\d+(_[A-Z0-9]+)*|F\d+|BF\d+|MXFP\d+(_[A-Z0-9]+)*)$/i; - -/** - * Matches prefix for custom quantization types, e.g. `UD-Q8_K_XL`. - */ -export const MODEL_CUSTOM_QUANTIZATION_PREFIX_RE = /^UD$/i; - -/** - * Matches a parameter-count segment, e.g. `7B`, `1.5b`, `120M`. - * The optional leading `E` covers effective-parameter sizes, e.g. Gemma's - * `E2B`/`E4B` (MatFormer models sized by resident params). - */ -export const MODEL_PARAMS_RE = /^[Ee]?\d+(\.\d+)?[BbMmKkTt]$/; - -/** - * Matches an activated-parameter-count segment, e.g. `A10B`, `a2.4b`. - * The leading `A`/`a` distinguishes it from a regular params segment. - */ -export const MODEL_ACTIVATED_PARAMS_RE = /^[Aa]\d+(\.\d+)?[BbMmKkTt]$/; - -/** - * Container format segments to exclude from tags (every model uses these). - */ -export const MODEL_IGNORED_SEGMENTS = new Set(['GGUF', 'GGML']); - -/** - * Matches a trailing weight file extension, e.g. `model.gguf` -> `model`. - */ -export const MODEL_WEIGHT_EXTENSION_RE = /\.(gguf|ggml)$/i; diff --git a/tools/ui/src/lib/constants/model-loading.ts b/tools/ui/src/lib/constants/model-loading.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/model-loading.ts rename to tools/ui/src/lib/constants/model-loading.constants.ts diff --git a/tools/ui/src/lib/constants/path-display.ts b/tools/ui/src/lib/constants/path-display.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/path-display.ts rename to tools/ui/src/lib/constants/path-display.constants.ts diff --git a/tools/ui/src/lib/constants/precision.ts b/tools/ui/src/lib/constants/precision.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/precision.ts rename to tools/ui/src/lib/constants/precision.constants.ts diff --git a/tools/ui/src/lib/constants/processing-info.ts b/tools/ui/src/lib/constants/processing-info.ts deleted file mode 100644 index 2c3f7dc53440..000000000000 --- a/tools/ui/src/lib/constants/processing-info.ts +++ /dev/null @@ -1,8 +0,0 @@ -export const PROCESSING_INFO_TIMEOUT = 2000; - -/** - * Statistics units labels - */ -export const STATS_UNITS = { - TOKENS_PER_SECOND: 't/s' -} as const; diff --git a/tools/ui/src/lib/constants/pwa.ts b/tools/ui/src/lib/constants/pwa.constants.ts similarity index 99% rename from tools/ui/src/lib/constants/pwa.ts rename to tools/ui/src/lib/constants/pwa.constants.ts index 4da37c06bcfe..e807f4a97a65 100644 --- a/tools/ui/src/lib/constants/pwa.ts +++ b/tools/ui/src/lib/constants/pwa.constants.ts @@ -3,7 +3,7 @@ * definitions across the codebase. */ -import { APP_NAME } from './app'; +import { APP_NAME } from './app.constants'; export const MEDIA_QUERIES = { DISPLAY_MODE_STANDALONE: '(display-mode: standalone)', diff --git a/tools/ui/src/lib/constants/reasoning-effort-tokens.ts b/tools/ui/src/lib/constants/reasoning-effort-tokens.ts deleted file mode 100644 index a0c244e5e9d3..000000000000 --- a/tools/ui/src/lib/constants/reasoning-effort-tokens.ts +++ /dev/null @@ -1,12 +0,0 @@ -import { ReasoningEffort } from '$lib/enums'; - -/** - * Reasoning effort to token budget mapping. - * Maps the ReasoningEffort enum values to concrete token counts for the server. - */ -export const REASONING_EFFORT_TOKENS: Record = { - [ReasoningEffort.HIGH]: 8192, - [ReasoningEffort.LOW]: 512, - [ReasoningEffort.MAX]: -1, // unlimited - [ReasoningEffort.MEDIUM]: 2048 -}; diff --git a/tools/ui/src/lib/constants/reasoning-effort.ts b/tools/ui/src/lib/constants/reasoning-effort.constants.ts similarity index 71% rename from tools/ui/src/lib/constants/reasoning-effort.ts rename to tools/ui/src/lib/constants/reasoning-effort.constants.ts index 4cbb7388ba19..e8ec5f0e8dd7 100644 --- a/tools/ui/src/lib/constants/reasoning-effort.ts +++ b/tools/ui/src/lib/constants/reasoning-effort.constants.ts @@ -22,3 +22,14 @@ export const REASONING_EFFORT_LEVELS: ReasoningEffortLevel[] = [ { label: 'High', value: ReasoningEffort.HIGH }, { hasInfo: true, label: 'Max', value: ReasoningEffort.MAX } ]; + +/** + * Reasoning effort to token budget mapping. + * Maps the ReasoningEffort enum values to concrete token counts for the server. + */ +export const REASONING_EFFORT_TOKENS: Record = { + [ReasoningEffort.HIGH]: 8192, + [ReasoningEffort.LOW]: 512, + [ReasoningEffort.MAX]: -1, // unlimited + [ReasoningEffort.MEDIUM]: 2048 +}; diff --git a/tools/ui/src/lib/constants/recommended-mcp-servers.ts b/tools/ui/src/lib/constants/recommended-mcp-servers.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/recommended-mcp-servers.ts rename to tools/ui/src/lib/constants/recommended-mcp-servers.constants.ts diff --git a/tools/ui/src/lib/constants/routes.ts b/tools/ui/src/lib/constants/routes.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/routes.ts rename to tools/ui/src/lib/constants/routes.constants.ts diff --git a/tools/ui/src/lib/constants/sandbox.constants.ts b/tools/ui/src/lib/constants/sandbox.constants.ts new file mode 100644 index 000000000000..9846e471a94e --- /dev/null +++ b/tools/ui/src/lib/constants/sandbox.constants.ts @@ -0,0 +1,13 @@ +import { BuiltInTool } from '$lib/enums'; + +export const SANDBOX_TOOL_NAME = BuiltInTool.RUN_JAVASCRIPT; + +export const SANDBOX_TIMEOUT_MS_DEFAULT = 10000; + +export const SANDBOX_TIMEOUT_MS_MAX = 30000; + +export const SANDBOX_OUTPUT_MAX_CHARS = 8192; + +export const SANDBOX_EMPTY_OUTPUT = '(no output)'; + +export const SANDBOX_TRUNCATION_NOTICE = '[output truncated]'; diff --git a/tools/ui/src/lib/constants/settings-keys.ts b/tools/ui/src/lib/constants/settings-keys.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/settings-keys.ts rename to tools/ui/src/lib/constants/settings-keys.constants.ts diff --git a/tools/ui/src/lib/constants/settings-registry.ts b/tools/ui/src/lib/constants/settings-registry.constants.ts similarity index 97% rename from tools/ui/src/lib/constants/settings-registry.ts rename to tools/ui/src/lib/constants/settings-registry.constants.ts index ada029a33ff5..e4445995702f 100644 --- a/tools/ui/src/lib/constants/settings-registry.ts +++ b/tools/ui/src/lib/constants/settings-registry.constants.ts @@ -1,12 +1,9 @@ -import { CLI_FLAGS } from './cli-flags'; -import { DEFAULT_MCP_CONFIG } from './mcp'; -import { ROUTES, SETTINGS_SECTION_SLUGS } from './routes'; -import { SETTINGS_KEYS } from './settings-keys'; -import { TITLE_GENERATION } from './title-generation'; -import { - FILE_GLOB_SEARCH_PICKERS_DEFAULT_SEARCH_DEPTH, - FILE_GLOB_SEARCH_PICKERS_MAX_SEARCH_DEPTH -} from './working-directory'; +import { CLI_FLAGS } from './cli-flags.constants'; +import { DEFAULT_MCP_CONFIG } from './mcp.constants'; +import { ROUTES, SETTINGS_SECTION_SLUGS } from './routes.constants'; +import { SETTINGS_KEYS } from './settings-keys.constants'; +import { TITLE_GENERATION } from './title-generation.constants'; +import { FILE_GLOB_SEARCH_PICKERS } from './working-directory.constants'; import { AlertTriangle, Code, @@ -14,7 +11,6 @@ import { Funnel, ListRestart, Monitor, - Monitor as MonitorIcon, Moon, PencilRuler, Sliders, @@ -53,7 +49,7 @@ const STANDALONE_SECTIONS: { title: SettingsSectionTitle; slug: string; icon: Co } ]; const COLOR_MODE_OPTIONS: Array<{ value: string; label: string; icon: Component }> = [ - { icon: MonitorIcon, label: 'System', value: ColorMode.SYSTEM }, + { icon: Monitor, label: 'System', value: ColorMode.SYSTEM }, { icon: Sun, label: 'Light', value: ColorMode.LIGHT }, { icon: Moon, label: 'Dark', value: ColorMode.DARK } ]; @@ -106,14 +102,14 @@ const SETTINGS_REGISTRY: Record = { type: SettingsFieldType.INPUT }, { - defaultValue: FILE_GLOB_SEARCH_PICKERS_DEFAULT_SEARCH_DEPTH, + defaultValue: FILE_GLOB_SEARCH_PICKERS.DEFAULT_SEARCH_DEPTH, help: 'How many directory levels below the working directory the @-mention file search descends. Larger values surface deeply nested files but take longer on large trees.', isPositiveInteger: true, key: SETTINGS_KEYS.MENTION_SEARCH_MAX_DEPTH, label: 'Mention search depth', - max: FILE_GLOB_SEARCH_PICKERS_MAX_SEARCH_DEPTH, + max: FILE_GLOB_SEARCH_PICKERS.MAX_SEARCH_DEPTH, min: 1, - placeholder: `${FILE_GLOB_SEARCH_PICKERS_DEFAULT_SEARCH_DEPTH}`, + placeholder: `${FILE_GLOB_SEARCH_PICKERS.DEFAULT_SEARCH_DEPTH}`, section: SETTINGS_SECTION_SLUGS.AGENTIC, type: SettingsFieldType.INPUT } diff --git a/tools/ui/src/lib/constants/special-characters.constants.ts b/tools/ui/src/lib/constants/special-characters.constants.ts new file mode 100644 index 000000000000..aaeebca33f99 --- /dev/null +++ b/tools/ui/src/lib/constants/special-characters.constants.ts @@ -0,0 +1,16 @@ +// Control / whitespace / formatting characters that appear literally inside rendered text. + +/** Line feed. */ +export const NEWLINE = '\n'; + +/** Horizontal tab. */ +export const TAB = '\t'; + +/** Non-breaking space. */ +export const NBSP = '\u00a0'; + +/** Non-breaking spaces used to render a tab stop that whitespace collapsing would otherwise squash. */ +export const TAB_AS_SPACES = NBSP.repeat(4); + +/** Matches a CR-terminated or bare LF line break. */ +export const LINE_BREAK = /\r?\n/; diff --git a/tools/ui/src/lib/constants/storage.ts b/tools/ui/src/lib/constants/storage.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/storage.ts rename to tools/ui/src/lib/constants/storage.constants.ts diff --git a/tools/ui/src/lib/constants/sse.ts b/tools/ui/src/lib/constants/stream.constants.ts similarity index 53% rename from tools/ui/src/lib/constants/sse.ts rename to tools/ui/src/lib/constants/stream.constants.ts index 0eb4b6edeea2..64f67243c266 100644 --- a/tools/ui/src/lib/constants/sse.ts +++ b/tools/ui/src/lib/constants/stream.constants.ts @@ -1,3 +1,11 @@ +// grace window after a visibilitychange before we kick a reader whose socket likely died +// while the tab was hidden. covers brief background pauses without thrashing live streams +export const STREAM_VISIBILITY_KICK_MS = 3000; + +// separator joining a conversation id and its per-model stream identity +// suffix (conv::model) used by the server side replay buffer +export const CONVERSATION_ID_SEPARATOR = '::'; + /** * Server-sent events wire format, shared by the chat stream and the * /models/sse status feed (text/event-stream). diff --git a/tools/ui/src/lib/constants/stream.ts b/tools/ui/src/lib/constants/stream.ts deleted file mode 100644 index 67951ee95301..000000000000 --- a/tools/ui/src/lib/constants/stream.ts +++ /dev/null @@ -1,3 +0,0 @@ -// grace window after a visibilitychange before we kick a reader whose socket likely died -// while the tab was hidden. covers brief background pauses without thrashing live streams -export const STREAM_VISIBILITY_KICK_MS = 3000; diff --git a/tools/ui/src/lib/constants/supported-file-types.ts b/tools/ui/src/lib/constants/supported-file-types.constants.ts similarity index 100% rename from tools/ui/src/lib/constants/supported-file-types.ts rename to tools/ui/src/lib/constants/supported-file-types.constants.ts diff --git a/tools/ui/src/lib/constants/svg-blocks.constants.ts b/tools/ui/src/lib/constants/svg-blocks.constants.ts new file mode 100644 index 000000000000..705800c2622c --- /dev/null +++ b/tools/ui/src/lib/constants/svg-blocks.constants.ts @@ -0,0 +1,57 @@ +/** + * Constants for rendering svg code blocks inline. + */ +export const SVG = { + // CSS classes applied to the inline svg block and its chrome. + BLOCK_CLASS: 'svg-block', + /** + * Shadow root style for the zoom dialog svg. Lets the svg grow past its + * intrinsic size so pan and zoom have room to work. + */ + DIALOG_SHADOW_STYLE: + ':host{display:inline-block}svg{min-height:min(50vh,12rem);min-width:min(80vw,20rem);max-width:none;max-height:none;height:auto;width:auto;display:block}', + ID_ATTR: 'data-svg-id', + + /** + * Shadow root style for an inline svg block. Mirrors the centered, padded + * sizing the light dom used before the svg moved behind a shadow boundary. + */ + INLINE_SHADOW_STYLE: + ':host{display:block;width:100%;text-align:center}svg{display:block;margin:0 auto;width:auto;height:auto;max-width:100%;max-height:70vh;min-height:8rem;padding:3rem 1rem}', + // Languages that mark a code block as svg content. + LANGUAGE: 'svg', + /** + * Hard size ceiling for a single inline svg block. + * Above this the source is left as raw text instead of being rendered. + */ + MAX_BYTES: 256 * 1024, + + RENDERED_ATTR: 'data-svg-rendered', + /** + * DOMPurify config for untrusted svg coming from model output. + * + * foreignObject and script stay forbidden unconditionally, they are the only + * inline svg vectors that execute arbitrary html or js. Everything else is + * allowed for maximum rendering compatibility: href and xlink:href stay so + * use, image, a and animateMotion work, and DOMPurify still neutralizes + * javascript: and data: uri schemes natively. External resource refs are + * allowed by design on a local first tool, the user browser fetches them. + * + * The sanitized svg is always mounted inside a shadow root (see svg-shadow), + * so an author {/if} diff --git a/tools/ui/src/routes/search/+page.svelte b/tools/ui/src/routes/search/+page.svelte index 424b74f67cf9..cb7f77f95e07 100644 --- a/tools/ui/src/routes/search/+page.svelte +++ b/tools/ui/src/routes/search/+page.svelte @@ -5,9 +5,7 @@ import { SearchInput, SidebarNavigationSearchResults } from '$lib/components/app'; import { ROUTES } from '$lib/constants'; import { RouterService } from '$lib/services/router.service'; - import { chatStore } from '$lib/stores/chat.svelte'; - import { conversations, conversationsStore } from '$lib/stores/conversations.svelte'; - import { isMobile } from '$lib/stores/viewport.svelte'; + import { chatStore, conversationsStore, isMobile } from '$lib/stores'; let searchQuery = $state(''); let searchInputRef = $state(null); @@ -19,7 +17,7 @@ if (query.length === 0) return []; - return conversations().filter((c) => c.name.toLowerCase().includes(query)); + return conversationsStore.conversations.filter((c) => c.name.toLowerCase().includes(query)); }); // Search page is intended for mobile; on desktop the sidebar already exposes @@ -35,7 +33,7 @@ } async function handleEditConversation(id: string) { - const conversation = conversations().find((c) => c.id === id); + const conversation = conversationsStore.conversations.find((c) => c.id === id); if (!conversation) return; @@ -47,7 +45,7 @@ } async function handleDeleteConversation(id: string) { - const conversation = conversations().find((c) => c.id === id); + const conversation = conversationsStore.conversations.find((c) => c.id === id); if (!conversation) return; diff --git a/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts b/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts index ddf1393e70e8..45af7e0d151c 100644 --- a/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts +++ b/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts @@ -1,7 +1,7 @@ import { CONFIG_LOCALSTORAGE_KEY, SETTING_CONFIG_DEFAULT } from '$lib/constants'; import { ParameterSyncService } from '$lib/services/parameter-sync.service'; import { serverStore } from '$lib/stores/server.svelte'; -import { config, settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings.svelte'; import type { SettingsConfigType } from '$lib/types'; import { beforeEach, describe, expect, it } from 'vitest'; @@ -40,7 +40,7 @@ function mockProps(uiSettings: Record) { const setUser = (key: string, value: Primitive) => settingsStore.updateConfig(key as keyof SettingsConfigType, value as never); -const current = (key: string) => (config() as Record)[key]; +const current = (key: string) => (settingsStore.config as Record)[key]; describe('registry-wide invariants', () => { beforeEach(() => { diff --git a/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts index 078cc7cbece3..32f4ff3dd4bd 100644 --- a/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts +++ b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts @@ -6,7 +6,7 @@ import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; import { MigrationService } from '$lib/services/migration.service'; -import { config, settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings.svelte'; import { beforeEach, describe, expect, it } from 'vitest'; const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1'; @@ -33,22 +33,22 @@ describe('renderContentAsRawText unfolding', () => { it('maps raw text to user content as plain text', async () => { await seedConfig({ renderContentAsRawText: true }); - expect(config().renderUserContentAsMarkdown).toBe(false); + expect(settingsStore.config.renderUserContentAsMarkdown).toBe(false); }); it('maps markdown to user content as markdown', async () => { await seedConfig({ renderContentAsRawText: false }); - expect(config().renderUserContentAsMarkdown).toBe(true); + expect(settingsStore.config.renderUserContentAsMarkdown).toBe(true); }); it('leaves thinking on its own default', async () => { await seedConfig({ renderContentAsRawText: true }); - expect(config().renderThinkingAsMarkdown).toBe(true); + expect(settingsStore.config.renderThinkingAsMarkdown).toBe(true); }); it('keeps an explicit user preference over the toggle', async () => { await seedConfig({ renderContentAsRawText: true, renderUserContentAsMarkdown: true }); - expect(config().renderUserContentAsMarkdown).toBe(true); + expect(settingsStore.config.renderUserContentAsMarkdown).toBe(true); }); it('drops the toggle from the persisted config', async () => { @@ -58,7 +58,7 @@ describe('renderContentAsRawText unfolding', () => { it('leaves both surfaces on markdown when nothing is stored', async () => { await seedConfig({}); - expect(config().renderUserContentAsMarkdown).toBe(true); - expect(config().renderThinkingAsMarkdown).toBe(true); + expect(settingsStore.config.renderUserContentAsMarkdown).toBe(true); + expect(settingsStore.config.renderThinkingAsMarkdown).toBe(true); }); }); diff --git a/tools/ui/tests/client/ui-settings-sync.svelte.test.ts b/tools/ui/tests/client/ui-settings-sync.svelte.test.ts index 110011ef0dcb..86e0898937d0 100644 --- a/tools/ui/tests/client/ui-settings-sync.svelte.test.ts +++ b/tools/ui/tests/client/ui-settings-sync.svelte.test.ts @@ -1,6 +1,6 @@ import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; import { serverStore } from '$lib/stores/server.svelte'; -import { config, settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings.svelte'; import { beforeEach, describe, expect, it } from 'vitest'; function mockProps(uiSettings: Record) { @@ -25,7 +25,7 @@ describe('server ui_settings application semantics', () => { settingsStore.syncWithServerDefaults(); - expect(config().theme).toBe('dark'); + expect(settingsStore.config.theme).toBe('dark'); }); it('never reapplies on later loads: the user config diverges freely', () => { @@ -40,8 +40,8 @@ describe('server ui_settings application semantics', () => { settingsStore.syncWithServerDefaults(); settingsStore.syncWithServerDefaults(); - expect(config().theme).toBe('light'); - expect(config().apiKey).toBe('sk-user-key'); + expect(settingsStore.config.theme).toBe('light'); + expect(settingsStore.config.apiKey).toBe('sk-user-key'); const stored = JSON.parse(localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}'); expect(stored.apiKey).toBe('sk-user-key'); @@ -55,8 +55,8 @@ describe('server ui_settings application semantics', () => { settingsStore.forceSyncWithServerDefaults(); - expect(config().theme).toBe('dark'); - expect(config().apiKey).toBe(''); + expect(settingsStore.config.theme).toBe('dark'); + expect(settingsStore.config.apiKey).toBe(''); }); }); From f2efd64141b9f9003c592375adf287f5fe1e5b80 Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Thu, 13 Aug 2026 08:13:52 +0200 Subject: [PATCH 048/211] ui: Move `styles/` to `$lib` scope (#26950) * refactor: Move `styles/` to `src/lib` and remove legacy alias * chore: Add newline --- tools/ui/src/app.html | 1 + .../app/content/MarkdownContent/MarkdownContent.svelte | 2 +- tools/ui/src/{ => lib}/styles/katex-custom.scss | 0 tools/ui/svelte.config.js | 3 --- 4 files changed, 2 insertions(+), 4 deletions(-) rename tools/ui/src/{ => lib}/styles/katex-custom.scss (100%) diff --git a/tools/ui/src/app.html b/tools/ui/src/app.html index e1de226dcb81..ef2787ad1bb4 100644 --- a/tools/ui/src/app.html +++ b/tools/ui/src/app.html @@ -2,6 +2,7 @@ + diff --git a/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte b/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte index 04d48f0aa860..cadd9f210baa 100644 --- a/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte @@ -1,5 +1,5 @@ {#if isMobile.current} - + {#snippet trigger({ disabled, onclick })} {/snippet} {:else} - + {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte index 13e97cc73c6f..d8fad772dd17 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte @@ -11,6 +11,7 @@ } from '$lib/components/app'; import { Button } from '$lib/components/ui/button'; import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants'; + import { setChatFormActionsContext } from '$lib/contexts'; import { FileTypeCategory, MessageRole } from '$lib/enums'; import { ChatService } from '$lib/services'; import { chatStore, conversationsStore, mcpStore, settingsStore } from '$lib/stores'; @@ -132,6 +133,42 @@ return livePromptTokens > 0 || liveOutputTokens > 0; }); + + setChatFormActionsContext({ + get disabled() { + return disabled; + }, + get hasAudioModality() { + return hasAudioModality; + }, + get hasMcpPromptsSupport() { + return hasMcpPromptsSupport; + }, + get hasMcpResourcesSupport() { + return hasMcpResourcesSupport; + }, + get hasVideoModality() { + return hasVideoModality; + }, + get hasVisionModality() { + return hasVisionModality; + }, + get onFileUpload() { + return onFileUpload; + }, + get onMcpPromptClick() { + return onMcpPromptClick; + }, + get onMcpResourcesClick() { + return onMcpResourcesClick; + }, + get onMcpSettingsClick() { + return () => goto(ROUTES.MCP_SERVERS); + }, + get onSystemPromptClick() { + return onSystemPromptClick; + } + });
{#if showAddButton}
- goto(ROUTES.MCP_SERVERS)} - /> +
{/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte index b3e54ada02fd..78cb8872173b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte @@ -8,16 +8,21 @@ ChatMessageUser } from '$lib/components/app/chat'; import { REASONING_TAGS, ROUTES, SYSTEM_MESSAGE_PLACEHOLDER } from '$lib/constants'; - import { getChatActionsContext, setMessageEditContext } from '$lib/contexts'; + import { setChatMessageActionsContext, setChatMessageEditContext } from '$lib/contexts'; import { AgenticSectionType, AttachmentType, MessageRole } from '$lib/enums'; import { DatabaseService } from '$lib/services/database.service'; import { chatStore, conversationsStore, isMobile } from '$lib/stores'; - import type { DatabaseMessageExtraMcpPrompt } from '$lib/types'; + import type { + ChatMessageActions, + ChatMessageDeletionInfo, + DatabaseMessageExtraMcpPrompt + } from '$lib/types'; import { deriveAgenticSections } from '$lib/utils'; import { parseFilesToMessageExtras } from '$lib/utils/browser-only'; interface Props { class?: string; + chatActions: ChatMessageActions; message: DatabaseMessage; toolMessages?: DatabaseMessage[]; isLastAssistantMessage?: boolean; @@ -27,6 +32,7 @@ } let { + chatActions, class: className = '', isLastAssistantMessage = false, isLastUserMessage = false, @@ -36,14 +42,7 @@ toolMessages = [] }: Props = $props(); - const chatActions = getChatActionsContext(); - - let deletionInfo = $state<{ - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - } | null>(null); + let deletionInfo = $state(null); // The system message placeholder must never surface as editable content; keeping // it in the derived (not just in handleEdit) guards against prop invalidation // reverting the override while editing @@ -112,7 +111,7 @@ let showSaveOnlyOption = $derived(message.role === MessageRole.USER); let showBranchAfterEditOption = $derived(message.role === MessageRole.ASSISTANT); - setMessageEditContext({ + setChatMessageEditContext({ cancel: handleCancelEdit, get editedContent() { return editedContent; @@ -166,6 +165,30 @@ startEdit: handleEdit }); + setChatMessageActionsContext({ + confirmDelete: handleConfirmDelete, + copy: handleCopy, + get deletionInfo() { + return deletionInfo; + }, + get forkConversation() { + const isForkableUser = message.role === MessageRole.USER && !mcpPromptExtra; + + return isForkableUser || message.role === MessageRole.ASSISTANT + ? handleForkConversation + : undefined; + }, + navigateToSibling: handleNavigateToSibling, + requestDelete: handleDelete, + setShowDeleteDialog: handleShowDeleteDialogChange, + get showDeleteDialog() { + return showDeleteDialog; + }, + get siblingInfo() { + return siblingInfo; + } + }); + let mcpPromptExtra = $derived.by(() => { if (message.role !== MessageRole.USER) return null; @@ -360,73 +383,22 @@
{#if message.role === MessageRole.SYSTEM} - + {:else if mcpPromptExtra} - + {:else if isSynthetic} {:else if message.role === MessageRole.USER} - + {:else} {/if}
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte index 14ea68099c6f..b92be9fbd61f 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte @@ -8,7 +8,7 @@ ChatMessageAssistantStatistics, ChatMessageEditForm } from '$lib/components/app'; - import { getMessageEditContext } from '$lib/contexts'; + import { getChatMessageEditContext } from '$lib/contexts'; import { MessageRole } from '$lib/enums'; import { useProcessingState } from '$lib/hooks/use-processing-state.svelte'; import { chatStore, modelsStore, serverStore, settingsStore } from '$lib/stores'; @@ -17,51 +17,26 @@ interface Props { class?: string; - deletionInfo: { - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - } | null; isLastAssistantMessage?: boolean; message: DatabaseMessage; toolMessages?: DatabaseMessage[]; - onCopy: () => void; - onConfirmDelete: () => void; onContinue?: () => void; - onDelete: () => void; - onEdit?: () => void; - onForkConversation?: (options: { name: string; includeAttachments: boolean }) => void; - onNavigateToSibling?: (siblingId: string) => void; onRegenerate: (modelOverride?: string) => void; - onShowDeleteDialogChange: (show: boolean) => void; - showDeleteDialog: boolean; - siblingInfo?: ChatMessageSiblingInfo | null; textareaElement?: HTMLTextAreaElement; } let { class: className = '', - deletionInfo, isLastAssistantMessage = false, message, - onConfirmDelete, onContinue, - onCopy, - onDelete, - onEdit, - onForkConversation, - onNavigateToSibling, onRegenerate, - onShowDeleteDialogChange, - showDeleteDialog, - siblingInfo = null, textareaElement = $bindable(), toolMessages = [] }: Props = $props(); // Get edit context - const editCtx = getMessageEditContext(); + const editCtx = getChatMessageEditContext(); const isAgentic = $derived(hasAgenticContent(message, toolMessages)); const processingState = useProcessingState(); @@ -207,18 +182,8 @@ role={MessageRole.ASSISTANT} justify="start" actionsPosition="left" - {siblingInfo} - {showDeleteDialog} - {deletionInfo} - {onCopy} - {onEdit} {onRegenerate} onContinue={currentConfig.enableContinueGeneration ? onContinue : undefined} - {onForkConversation} - {onDelete} - {onConfirmDelete} - {onNavigateToSibling} - {onShowDeleteDialogChange} showRawOutputSwitch={currentConfig.showRawOutputSwitch} rawOutputEnabled={showRawOutput} onRawOutputToggle={(enabled) => (showRawOutput = enabled)} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte index d3aa7251e096..4563b1fa8630 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageMcpPrompt/ChatMessageMcpPrompt.svelte @@ -4,7 +4,7 @@ ChatMessageEditForm, ChatMessageMcpPromptContent } from '$lib/components/app'; - import { getMessageEditContext } from '$lib/contexts'; + import { getChatMessageEditContext } from '$lib/contexts'; import { McpPromptVariant, MessageRole } from '$lib/enums'; import type { DatabaseMessageExtraMcpPrompt } from '$lib/types'; @@ -12,39 +12,12 @@ class?: string; message: DatabaseMessage; mcpPrompt: DatabaseMessageExtraMcpPrompt; - siblingInfo?: ChatMessageSiblingInfo | null; - showDeleteDialog: boolean; - deletionInfo: { - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - } | null; - onCopy: () => void; - onEdit: () => void; - onDelete: () => void; - onConfirmDelete: () => void; - onNavigateToSibling?: (siblingId: string) => void; - onShowDeleteDialogChange: (show: boolean) => void; } - let { - class: className = '', - deletionInfo, - mcpPrompt, - message, - onConfirmDelete, - onCopy, - onDelete, - onEdit, - onNavigateToSibling, - onShowDeleteDialogChange, - showDeleteDialog, - siblingInfo = null - }: Props = $props(); + let { class: className = '', mcpPrompt, message }: Props = $props(); // Get edit context - const editCtx = getMessageEditContext(); + const editCtx = getChatMessageEditContext();
- +
{/if} {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte index 2f919f9eb89c..c6222f568c48 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageSystem/ChatMessageSystem.svelte @@ -4,7 +4,7 @@ import { Button } from '$lib/components/ui/button'; import { Card } from '$lib/components/ui/card'; import { INPUT_CLASSES } from '$lib/constants'; - import { getMessageEditContext } from '$lib/contexts'; + import { getChatMessageEditContext } from '$lib/contexts'; import { KeyboardKey, MessageRole } from '$lib/enums'; import { settingsStore } from '$lib/stores'; import { autoResizeTextarea, isIMEComposing } from '$lib/utils'; @@ -12,39 +12,12 @@ interface Props { class?: string; message: DatabaseMessage; - siblingInfo?: ChatMessageSiblingInfo | null; - showDeleteDialog: boolean; - deletionInfo: { - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - } | null; - onCopy: () => void; - onEdit: () => void; - onDelete: () => void; - onConfirmDelete: () => void; - onNavigateToSibling?: (siblingId: string) => void; - onShowDeleteDialogChange: (show: boolean) => void; textareaElement?: HTMLTextAreaElement; } - let { - class: className = '', - deletionInfo, - message, - onConfirmDelete, - onCopy, - onDelete, - onEdit, - onNavigateToSibling, - onShowDeleteDialogChange, - showDeleteDialog, - siblingInfo = null, - textareaElement = $bindable() - }: Props = $props(); - - const editCtx = getMessageEditContext(); + let { class: className = '', message, textareaElement = $bindable() }: Props = $props(); + + const editCtx = getChatMessageEditContext(); function handleEditKeydown(event: KeyboardEvent) { if (event.key === KeyboardKey.ENTER && !event.shiftKey && !isIMEComposing(event)) { @@ -218,20 +191,7 @@ {#if message.timestamp}
- +
{/if} {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte index db09bb7235d3..a6fa2e250471 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte @@ -12,8 +12,7 @@ import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte'; import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte'; import { BuiltInTool } from '$lib/enums'; - import type { AgenticSection } from '$lib/types'; - import type { DatabaseMessageExtra } from '$lib/types'; + import type { AgenticSection, DatabaseMessageExtra } from '$lib/types'; import { extractSearchQuery, extractSearchResults, isWebSearchToolName } from '$lib/utils'; interface Props { diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte index 639bee89784a..d7103bf97f27 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockExecShellCommand.svelte @@ -13,8 +13,7 @@ import { SETTINGS_KEYS, TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants'; import { AttachmentType } from '$lib/enums'; import { settingsStore, toolsStore } from '$lib/stores'; - import type { AgenticSection, ToolResultLine } from '$lib/types'; - import type { DatabaseMessageExtra } from '$lib/types'; + import type { AgenticSection, DatabaseMessageExtra, ToolResultLine } from '$lib/types'; import { abbreviateHome, type ExecShellExitStatus, diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte index f902ec88d594..d217a9c382ee 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte @@ -5,7 +5,7 @@ ChatMessageStatistics, ChatMessageUserBubble } from '$lib/components/app/chat'; - import { getMessageEditContext } from '$lib/contexts'; + import { getChatMessageEditContext } from '$lib/contexts'; import { ChatMessageStatisticsMode, MessageRole } from '$lib/enums'; import { useProcessingState } from '$lib/hooks/use-processing-state.svelte'; import { chatStore, settingsStore } from '$lib/stores'; @@ -13,44 +13,19 @@ interface Props { class?: string; message: DatabaseMessage; - siblingInfo?: ChatMessageSiblingInfo | null; - deletionInfo: { - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - } | null; isLastUserMessage?: boolean; nextAssistantMessage?: DatabaseMessage | null; - showDeleteDialog: boolean; - onEdit: () => void; - onDelete: () => void; - onConfirmDelete: () => void; - onForkConversation?: (options: { name: string; includeAttachments: boolean }) => void; - onShowDeleteDialogChange: (show: boolean) => void; - onNavigateToSibling?: (siblingId: string) => void; - onCopy: () => void; } let { class: className = '', - deletionInfo, isLastUserMessage = false, message, - nextAssistantMessage = null, - onConfirmDelete, - onCopy, - onDelete, - onEdit, - onForkConversation, - onNavigateToSibling, - onShowDeleteDialogChange, - showDeleteDialog, - siblingInfo = null + nextAssistantMessage = null }: Props = $props(); // Get contexts - const editCtx = getMessageEditContext(); + const editCtx = getChatMessageEditContext(); const processingState = useProcessingState(); const currentConfig = $derived(settingsStore.config); @@ -132,21 +107,7 @@ {#if message.timestamp}
- +
{/if} {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte index f3e8c4e91d39..a072f2e84dd2 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte @@ -1,7 +1,7 @@ @@ -88,18 +70,16 @@ ? 'left-0' : 'right-0'} flex items-center gap-2 opacity-100 transition-opacity" > - {#if siblingInfo && siblingInfo.totalSiblings > 1} - + {#if messageActions.siblingInfo && messageActions.siblingInfo.totalSiblings > 1} + {/if}
- + - {#if onEdit} - - {/if} + {#if role === MessageRole.ASSISTANT && onRegenerate} onRegenerate()} /> @@ -109,11 +89,11 @@ {/if} - {#if onForkConversation} + {#if messageActions.forkConversation} {/if} - +
@@ -129,19 +109,19 @@ 1 - ? `This will delete ${deletionInfo.totalCount} messages including: ${deletionInfo.userMessages} user message${deletionInfo.userMessages > 1 ? 's' : ''} and ${deletionInfo.assistantMessages} assistant response${deletionInfo.assistantMessages > 1 ? 's' : ''}. All messages in this branch and their responses will be permanently removed. This action cannot be undone.` + description={messageActions.deletionInfo && messageActions.deletionInfo.totalCount > 1 + ? `This will delete ${messageActions.deletionInfo.totalCount} messages including: ${messageActions.deletionInfo.userMessages} user message${messageActions.deletionInfo.userMessages > 1 ? 's' : ''} and ${messageActions.deletionInfo.assistantMessages} assistant response${messageActions.deletionInfo.assistantMessages > 1 ? 's' : ''}. All messages in this branch and their responses will be permanently removed. This action cannot be undone.` : 'Are you sure you want to delete this message? This action cannot be undone.'} - confirmText={deletionInfo && deletionInfo.totalCount > 1 - ? `Delete ${deletionInfo.totalCount} Messages` + confirmText={messageActions.deletionInfo && messageActions.deletionInfo.totalCount > 1 + ? `Delete ${messageActions.deletionInfo.totalCount} Messages` : 'Delete'} cancelText="Cancel" variant="destructive" icon={Trash2} onConfirm={handleConfirmDelete} - onCancel={() => onShowDeleteDialogChange(false)} + onCancel={() => messageActions.setShowDeleteDialog(false)} /> import { ChevronLeft, ChevronRight } from '@lucide/svelte'; import { ActionIcon } from '$lib/components/app'; + import { getChatMessageActionsContext } from '$lib/contexts'; interface Props { class?: string; - siblingInfo: ChatMessageSiblingInfo | null; - onNavigateToSibling?: (siblingId: string) => void; } - let { class: className = '', onNavigateToSibling, siblingInfo }: Props = $props(); + let { class: className = '' }: Props = $props(); + + const messageActions = getChatMessageActionsContext(); + + let siblingInfo = $derived(messageActions.siblingInfo); let hasPrevious = $derived(siblingInfo && siblingInfo.currentIndex > 0); let hasNext = $derived(siblingInfo && siblingInfo.currentIndex < siblingInfo.totalSiblings - 1); @@ -31,7 +34,7 @@ tooltip="Previous version" disabled={!hasPrevious} class="h-5 w-5 p-0 {!hasPrevious ? '!cursor-not-allowed opacity-30' : ''}" - onclick={() => onNavigateToSibling?.(previousSiblingId!)} + onclick={() => messageActions.navigateToSibling(previousSiblingId!)} /> @@ -43,7 +46,7 @@ tooltip="Next version" disabled={!hasNext} class="h-5 w-5 p-0 {!hasNext ? 'opacity-30' : ''}" - onclick={() => onNavigateToSibling?.(nextSiblingId!)} + onclick={() => messageActions.navigateToSibling(nextSiblingId!)} /> {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte index 8ce1ead01e9f..5849799794f5 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte @@ -9,8 +9,8 @@ } from '$lib/components/app'; import { AgenticSectionType, ChatMessageStatsView, ToolPermissionDecision } from '$lib/enums'; import { agenticStore, settingsStore } from '$lib/stores'; - import type { AgenticSection } from '$lib/types'; import type { + AgenticSection, ChatMessageAgenticTimings, ChatMessageAgenticTurnStats, DatabaseMessage diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte index feacd0ddde34..369b6137b423 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte @@ -3,12 +3,12 @@ import { ChatForm, DialogConfirmation } from '$lib/components/app'; import { Button } from '$lib/components/ui/button'; import { Switch } from '$lib/components/ui/switch'; - import { getMessageEditContext } from '$lib/contexts'; + import { getChatMessageEditContext } from '$lib/contexts'; import { KeyboardKey, MessageRole } from '$lib/enums'; import { chatStore } from '$lib/stores'; import { processFilesToChatUploaded } from '$lib/utils/browser-only'; - const editCtx = getMessageEditContext(); + const editCtx = getChatMessageEditContext(); let saveWithoutRegenerate = $state(false); let showDiscardDialog = $state(false); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte index 78c9dcb25a74..a0861fb9619f 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageReasoningBlock.svelte @@ -4,8 +4,7 @@ import { REASONING_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants'; import { AgenticSectionType } from '$lib/enums'; import { settingsStore } from '$lib/stores'; - import type { DatabaseMessageExtra } from '$lib/types'; - import type { AgenticSection } from '$lib/types'; + import type { AgenticSection, DatabaseMessageExtra } from '$lib/types'; interface Props { section: AgenticSection; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte index 33d971f93188..2a8f45ba5388 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte @@ -1,8 +1,8 @@
diff --git a/tools/ui/src/lib/constants/context-keys.constants.ts b/tools/ui/src/lib/constants/context-keys.constants.ts index 0bd733b37066..62ff5413e134 100644 --- a/tools/ui/src/lib/constants/context-keys.constants.ts +++ b/tools/ui/src/lib/constants/context-keys.constants.ts @@ -1,3 +1,3 @@ -export const CONTEXT_KEY_MESSAGE_EDIT = 'chat-message-edit'; -export const CONTEXT_KEY_CHAT_ACTIONS = 'chat-actions'; -export const CONTEXT_KEY_CHAT_SETTINGS_CONFIG = 'chat-settings-config'; +export const CONTEXT_KEY_CHAT_MESSAGE_EDIT = 'chat-message-edit'; +export const CONTEXT_KEY_CHAT_MESSAGE_ACTIONS = 'chat-message-actions'; +export const CONTEXT_KEY_CHAT_FORM_ACTIONS = 'chat-form-actions'; diff --git a/tools/ui/src/lib/contexts/chat-actions.context.ts b/tools/ui/src/lib/contexts/chat-actions.context.ts deleted file mode 100644 index dffb5f3a36d9..000000000000 --- a/tools/ui/src/lib/contexts/chat-actions.context.ts +++ /dev/null @@ -1,39 +0,0 @@ -import { CONTEXT_KEY_CHAT_ACTIONS } from '$lib/constants'; -import { getContext, setContext } from 'svelte'; - -export interface ChatActionsContext { - copy: (message: DatabaseMessage) => void; - delete: (message: DatabaseMessage) => void; - navigateToSibling: (siblingId: string) => void; - editWithBranching: ( - message: DatabaseMessage, - newContent: string, - newExtras?: DatabaseMessageExtra[] - ) => void; - editWithReplacement: ( - message: DatabaseMessage, - newContent: string, - shouldBranch: boolean - ) => void; - editUserMessagePreserveResponses: ( - message: DatabaseMessage, - newContent: string, - newExtras?: DatabaseMessageExtra[] - ) => void; - regenerateWithBranching: (message: DatabaseMessage, modelOverride?: string) => void; - continueAssistantMessage: (message: DatabaseMessage) => void; - forkConversation: ( - message: DatabaseMessage, - options: { name: string; includeAttachments: boolean } - ) => void; -} - -const CHAT_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_ACTIONS); - -export function setChatActionsContext(ctx: ChatActionsContext): ChatActionsContext { - return setContext(CHAT_ACTIONS_KEY, ctx); -} - -export function getChatActionsContext(): ChatActionsContext { - return getContext(CHAT_ACTIONS_KEY); -} diff --git a/tools/ui/src/lib/contexts/chat-form-actions.context.ts b/tools/ui/src/lib/contexts/chat-form-actions.context.ts new file mode 100644 index 000000000000..a49f17447e44 --- /dev/null +++ b/tools/ui/src/lib/contexts/chat-form-actions.context.ts @@ -0,0 +1,19 @@ +import { CONTEXT_KEY_CHAT_FORM_ACTIONS } from '$lib/constants'; +import type { ChatFormActionsContext } from '$lib/types'; +import { getContext, setContext } from 'svelte'; + +const CHAT_FORM_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_FORM_ACTIONS); + +/** + * Sets the chat form actions context. Call in the parent component (ChatFormActions.svelte). + */ +export function setChatFormActionsContext(ctx: ChatFormActionsContext): ChatFormActionsContext { + return setContext(CHAT_FORM_ACTIONS_KEY, ctx); +} + +/** + * Gets the chat form actions context. Call in child components. + */ +export function getChatFormActionsContext(): ChatFormActionsContext { + return getContext(CHAT_FORM_ACTIONS_KEY); +} diff --git a/tools/ui/src/lib/contexts/chat-message-actions.context.ts b/tools/ui/src/lib/contexts/chat-message-actions.context.ts new file mode 100644 index 000000000000..fb075b3b0281 --- /dev/null +++ b/tools/ui/src/lib/contexts/chat-message-actions.context.ts @@ -0,0 +1,21 @@ +import { CONTEXT_KEY_CHAT_MESSAGE_ACTIONS } from '$lib/constants'; +import type { ChatMessageActionsContext } from '$lib/types'; +import { getContext, setContext } from 'svelte'; + +const CHAT_MESSAGE_ACTIONS_KEY = Symbol.for(CONTEXT_KEY_CHAT_MESSAGE_ACTIONS); + +/** + * Sets the per-message actions context. Call this in the parent component (ChatMessage.svelte). + */ +export function setChatMessageActionsContext( + ctx: ChatMessageActionsContext +): ChatMessageActionsContext { + return setContext(CHAT_MESSAGE_ACTIONS_KEY, ctx); +} + +/** + * Gets the per-message actions context. Call this in child components. + */ +export function getChatMessageActionsContext(): ChatMessageActionsContext { + return getContext(CHAT_MESSAGE_ACTIONS_KEY); +} diff --git a/tools/ui/src/lib/contexts/chat-message-edit.context.ts b/tools/ui/src/lib/contexts/chat-message-edit.context.ts new file mode 100644 index 000000000000..e9c053036e79 --- /dev/null +++ b/tools/ui/src/lib/contexts/chat-message-edit.context.ts @@ -0,0 +1,19 @@ +import { CONTEXT_KEY_CHAT_MESSAGE_EDIT } from '$lib/constants'; +import type { ChatMessageEditContext } from '$lib/types'; +import { getContext, setContext } from 'svelte'; + +const CHAT_MESSAGE_EDIT_KEY = Symbol.for(CONTEXT_KEY_CHAT_MESSAGE_EDIT); + +/** + * Sets the message edit context. Call this in the parent component (ChatMessage.svelte). + */ +export function setChatMessageEditContext(ctx: ChatMessageEditContext): ChatMessageEditContext { + return setContext(CHAT_MESSAGE_EDIT_KEY, ctx); +} + +/** + * Gets the message edit context. Call this in child components. + */ +export function getChatMessageEditContext(): ChatMessageEditContext { + return getContext(CHAT_MESSAGE_EDIT_KEY); +} diff --git a/tools/ui/src/lib/contexts/chat-settings-config.context.ts b/tools/ui/src/lib/contexts/chat-settings-config.context.ts deleted file mode 100644 index a90f709b49f8..000000000000 --- a/tools/ui/src/lib/contexts/chat-settings-config.context.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { CONTEXT_KEY_CHAT_SETTINGS_CONFIG } from '$lib/constants'; -import { getContext, setContext } from 'svelte'; - -export interface ChatSettingsConfigContext { - readonly localConfig: SettingsConfigType; - handleConfigChange: (key: string, value: string | boolean) => void; - handleThemeChange: (theme: string) => void; -} - -const CHAT_SETTINGS_CONFIG_KEY = Symbol.for(CONTEXT_KEY_CHAT_SETTINGS_CONFIG); - -export function setChatSettingsConfigContext( - ctx: ChatSettingsConfigContext -): ChatSettingsConfigContext { - return setContext(CHAT_SETTINGS_CONFIG_KEY, ctx); -} - -export function getChatSettingsConfigContext(): ChatSettingsConfigContext { - return getContext(CHAT_SETTINGS_CONFIG_KEY); -} diff --git a/tools/ui/src/lib/contexts/index.ts b/tools/ui/src/lib/contexts/index.ts index c6719fa9e477..4aaec8148c29 100644 --- a/tools/ui/src/lib/contexts/index.ts +++ b/tools/ui/src/lib/contexts/index.ts @@ -1,19 +1,8 @@ -export { - getMessageEditContext, - setMessageEditContext, - type MessageEditContext, - type MessageEditState, - type MessageEditActions -} from './message-edit.context'; +export { getChatMessageEditContext, setChatMessageEditContext } from './chat-message-edit.context'; export { - getChatActionsContext, - setChatActionsContext, - type ChatActionsContext -} from './chat-actions.context'; + getChatMessageActionsContext, + setChatMessageActionsContext +} from './chat-message-actions.context'; -export { - getChatSettingsConfigContext, - setChatSettingsConfigContext, - type ChatSettingsConfigContext -} from './chat-settings-config.context'; +export { getChatFormActionsContext, setChatFormActionsContext } from './chat-form-actions.context'; diff --git a/tools/ui/src/lib/contexts/message-edit.context.ts b/tools/ui/src/lib/contexts/message-edit.context.ts deleted file mode 100644 index 80f3e6eee28c..000000000000 --- a/tools/ui/src/lib/contexts/message-edit.context.ts +++ /dev/null @@ -1,51 +0,0 @@ -import { CONTEXT_KEY_MESSAGE_EDIT } from '$lib/constants'; -import { MessageRole } from '$lib/enums'; -import { getContext, setContext } from 'svelte'; - -export interface MessageEditState { - readonly isEditing: boolean; - readonly editedContent: string; - readonly editedExtras: DatabaseMessageExtra[]; - readonly editedUploadedFiles: ChatUploadedFile[]; - readonly originalContent: string; - readonly originalExtras: DatabaseMessageExtra[]; - readonly showSaveOnlyOption: boolean; - readonly showBranchAfterEditOption: boolean; - readonly shouldBranchAfterEdit: boolean; - readonly messageRole: MessageRole; - readonly rawEditContent?: string; -} - -export interface MessageEditActions { - setContent: (content: string) => void; - setExtras: (extras: DatabaseMessageExtra[]) => void; - setUploadedFiles: (files: ChatUploadedFile[]) => void; - save: () => void; - saveOnly: () => void; - cancel: () => void; - startEdit: () => void; -} - -export interface AssistantEditActions { - setShouldBranchAfterEdit: (value: boolean) => void; -} - -export type MessageEditContext = MessageEditState & - MessageEditActions & - Partial; - -const MESSAGE_EDIT_KEY = Symbol.for(CONTEXT_KEY_MESSAGE_EDIT); - -/** - * Sets the message edit context. Call this in the parent component (ChatMessage.svelte). - */ -export function setMessageEditContext(ctx: MessageEditContext): MessageEditContext { - return setContext(MESSAGE_EDIT_KEY, ctx); -} - -/** - * Gets the message edit context. Call this in child components. - */ -export function getMessageEditContext(): MessageEditContext { - return getContext(MESSAGE_EDIT_KEY); -} diff --git a/tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts b/tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts similarity index 91% rename from tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts rename to tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts index 271675404579..de2994e739a5 100644 --- a/tools/ui/src/lib/hooks/use-message-edit-context.svelte.ts +++ b/tools/ui/src/lib/hooks/use-chat-message-edit-context.svelte.ts @@ -1,15 +1,15 @@ -import { setMessageEditContext } from '$lib/contexts'; +import { setChatMessageEditContext } from '$lib/contexts'; import { MessageRole } from '$lib/enums'; import { parseFilesToMessageExtras } from '$lib/utils/convert-files-to-extra'; -interface UseMessageEditContextOptions { +interface UseChatMessageEditContextOptions { getContent: () => string; getExtras: () => DatabaseMessageExtra[]; showSaveOnlyOption?: boolean; onSave: (content: string, extras?: DatabaseMessageExtra[]) => void; } -export function useMessageEditContext(options: UseMessageEditContextOptions) { +export function useChatMessageEditContext(options: UseChatMessageEditContextOptions) { let isEditing = $state(false); let editedContent = $state(''); let editedExtras = $state([]); @@ -45,7 +45,7 @@ export function useMessageEditContext(options: UseMessageEditContextOptions) { isEditing = false; } - setMessageEditContext({ + setChatMessageEditContext({ cancel: handleCancelEdit, get editedContent() { return editedContent; diff --git a/tools/ui/src/lib/types/chat.d.ts b/tools/ui/src/lib/types/chat.d.ts index cb93e7795150..f0f3a297e8e1 100644 --- a/tools/ui/src/lib/types/chat.d.ts +++ b/tools/ui/src/lib/types/chat.d.ts @@ -7,7 +7,8 @@ import type { AttachmentMenuItemId, ChatFormCommandAction, ErrorDialogType, - FileMentionEntryType + FileMentionEntryType, + MessageRole } from '$lib/enums'; import type { Component } from 'svelte'; @@ -235,3 +236,111 @@ export interface ChatFormCommand { action: ChatFormCommandAction; disabled: boolean; } + +/** + * Data shown in the message delete confirmation dialog. + */ +export interface ChatMessageDeletionInfo { + totalCount: number; + userMessages: number; + assistantMessages: number; + messageTypes: string[]; +} + +/** + * Conversation-level message operations owned by ChatMessages (store calls + list + * refresh + user-action notification), passed to each ChatMessage as a prop. + */ +export interface ChatMessageActions { + copy: (message: DatabaseMessage) => void; + delete: (message: DatabaseMessage) => void; + navigateToSibling: (siblingId: string) => void; + editWithBranching: ( + message: DatabaseMessage, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ) => void; + editWithReplacement: ( + message: DatabaseMessage, + newContent: string, + shouldBranch: boolean + ) => void; + editUserMessagePreserveResponses: ( + message: DatabaseMessage, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ) => void; + regenerateWithBranching: (message: DatabaseMessage, modelOverride?: string) => void; + continueAssistantMessage: (message: DatabaseMessage) => void; + forkConversation: ( + message: DatabaseMessage, + options: { name: string; includeAttachments: boolean } + ) => void; +} + +/** + * Per-message actions and state. Set once per message in ChatMessage.svelte and + * consumed by its descendants (action icons, branching controls). + */ +export interface ChatMessageActionsContext { + readonly siblingInfo: ChatMessageSiblingInfo | null; + readonly deletionInfo: ChatMessageDeletionInfo | null; + readonly showDeleteDialog: boolean; + copy: () => void; + requestDelete: () => void; + confirmDelete: () => void; + setShowDeleteDialog: (show: boolean) => void; + navigateToSibling: (siblingId: string) => void; + forkConversation?: (options: { name: string; includeAttachments: boolean }) => void; +} + +export interface ChatMessageEditState { + readonly isEditing: boolean; + readonly editedContent: string; + readonly editedExtras: DatabaseMessageExtra[]; + readonly editedUploadedFiles: ChatUploadedFile[]; + readonly originalContent: string; + readonly originalExtras: DatabaseMessageExtra[]; + readonly showSaveOnlyOption: boolean; + readonly showBranchAfterEditOption: boolean; + readonly shouldBranchAfterEdit: boolean; + readonly messageRole: MessageRole; + readonly rawEditContent?: string; +} + +export interface ChatMessageEditActions { + setContent: (content: string) => void; + setExtras: (extras: DatabaseMessageExtra[]) => void; + setUploadedFiles: (files: ChatUploadedFile[]) => void; + save: () => void; + saveOnly: () => void; + cancel: () => void; + startEdit: () => void; +} + +export interface ChatMessageAssistantEditActions { + setShouldBranchAfterEdit: (value: boolean) => void; +} + +export type ChatMessageEditContext = ChatMessageEditState & + ChatMessageEditActions & + Partial; + +/** + * Actions and capability flags for the ChatForm add-menu. Set once in + * ChatFormActions.svelte and consumed by its deep descendants (the add sheet, + * dropdown and MCP servers submenu) to avoid relaying them through props. + */ +export interface ChatFormActionsContext { + readonly disabled: boolean; + readonly hasAudioModality: boolean; + readonly hasVideoModality: boolean; + readonly hasVisionModality: boolean; + readonly hasMcpPromptsSupport: boolean; + readonly hasMcpResourcesSupport: boolean; + onFileUpload?: () => void; + onSystemPromptClick?: () => void; + onMcpPromptClick?: () => void; + onMcpResourcesClick?: () => void; + onMcpSettingsClick?: () => void; +} diff --git a/tools/ui/src/lib/types/index.ts b/tools/ui/src/lib/types/index.ts index 917c9516b578..34767e2ae7df 100644 --- a/tools/ui/src/lib/types/index.ts +++ b/tools/ui/src/lib/types/index.ts @@ -44,6 +44,14 @@ export type { ChatUploadedFile, ChatAttachmentDisplayItem, ChatMessageSiblingInfo, + ChatMessageActions, + ChatMessageActionsContext, + ChatMessageDeletionInfo, + ChatMessageEditContext, + ChatMessageEditState, + ChatMessageEditActions, + ChatMessageAssistantEditActions, + ChatFormActionsContext, ChatMessagePromptProgress, ChatMessageTimings, ChatMessageAgenticTimings, diff --git a/tools/ui/tests/stories/ChatMessage.stories.svelte b/tools/ui/tests/stories/ChatMessage.stories.svelte index 023ba1ac4fc7..84fee2ea1c78 100644 --- a/tools/ui/tests/stories/ChatMessage.stories.svelte +++ b/tools/ui/tests/stories/ChatMessage.stories.svelte @@ -1,6 +1,7 @@ - +
- {#if useContenteditable} - { - pickers.handleInput(); - onValueChange?.(value); - }} - onPaste={handlePaste} - {disabled} - {placeholder} - /> - {:else} - { - pickers.handleInput(); - onValueChange?.(value); - }} - onPaste={handlePaste} - {disabled} - {placeholder} - /> - {/if} + { + pickers.handleInput(); + onValueChange?.(value); + }} + onPaste={handlePaste} + {disabled} + {placeholder} + {useContenteditable} + /> {#if mcpResourceStore.hasAttachments} {#if toolsStore.hasEnabledCwdTools} - - import ChatFormWorkingDirectoryChip from './ChatFormWorkingDirectoryChip.svelte'; - import ChatFormWorkingDirectoryResultsList from './ChatFormWorkingDirectoryResultsList.svelte'; + import ChatFormCurrentWorkingDirectoryChip from './ChatFormCurrentWorkingDirectoryChip.svelte'; + import ChatFormCurrentWorkingDirectoryResultsList from './ChatFormCurrentWorkingDirectoryResultsList.svelte'; import { FolderOpen } from '@lucide/svelte'; import SearchInput from '$lib/components/app/forms/SearchInput.svelte'; import * as Popover from '$lib/components/ui/popover'; @@ -250,7 +250,7 @@ // user cancelled - silently ignore; other errors are logged if (err instanceof DOMException && err.name === 'AbortError') return; - console.error('[ChatFormWorkingDirectory] showDirectoryPicker failed:', err); + console.error('[ChatFormCurrentWorkingDirectory] showDirectoryPicker failed:', err); } } @@ -331,7 +331,7 @@ onclick={onOpen} {disabled} > - {searchUnavailableMessage}
{:else if query.trim() && (search.isSearching || queryResults.length > 0 || searchError)} - + import ChatFormInputBasic from './ChatFormInputBasic.svelte'; + import ChatFormInputRich from './ChatFormInputRich.svelte'; + + interface Props { + class?: string; + disabled?: boolean; + onInput?: () => void; + onKeydown?: (event: KeyboardEvent) => void; + onPaste?: (event: ClipboardEvent) => void; + placeholder?: string; + value?: string; + useContenteditable?: boolean; + } + + let { + class: className = '', + disabled = false, + onInput, + onKeydown, + onPaste, + placeholder = 'Ask anything...', + useContenteditable = false, + value = $bindable('') + }: Props = $props(); + + let basicRef: ChatFormInputBasic | undefined = $state(); + let richRef: ChatFormInputRich | undefined = $state(); + + // The two renderers share one imperative handle (focus/caret/height), so + // the parent can drive whichever variant is mounted through this one. + export function getElement() { + return useContenteditable ? richRef?.getElement() : basicRef?.getElement(); + } + + export function focus() { + if (useContenteditable) richRef?.focus(); + else basicRef?.focus(); + } + + export function resetHeight() { + if (useContenteditable) richRef?.resetHeight(); + else basicRef?.resetHeight(); + } + + export function getCaretOffset(): number { + return useContenteditable + ? (richRef?.getCaretOffset() ?? 0) + : (basicRef?.getCaretOffset() ?? 0); + } + + export function setCaretOffset(offset: number) { + if (useContenteditable) richRef?.setCaretOffset(offset); + else basicRef?.setCaretOffset(offset); + } + + +{#if useContenteditable} + +{:else} + +{/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormTextarea.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputBasic.svelte similarity index 100% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormTextarea.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputBasic.svelte diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormFileInputInvisible.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputFileInputInvisible.svelte similarity index 100% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormFileInputInvisible.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputFileInputInvisible.svelte diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContentEditable.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte similarity index 97% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContentEditable.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte index ba44304011d9..f25c3c95cd2e 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContentEditable.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte @@ -2,7 +2,7 @@ import { CODE_BLOCK } from '$lib/constants'; import { ColorMode } from '$lib/enums'; import { isMobile } from '$lib/stores'; - import type { ContentEditableToken } from '$lib/types'; + import type { ChatFormInputRichToken } from '$lib/types'; import type { SourceHistoryEntry } from '$lib/utils'; import { badgeAwareWordJump, @@ -64,7 +64,7 @@ rootElement.dataset.empty = source.length === 0 ? 'true' : 'false'; } - function renderTokens(tokens: ContentEditableToken[]) { + function renderTokens(tokens: ChatFormInputRichToken[]) { if (!rootElement) return; const caret = rangeToTextOffset(rootElement, safeRange()); @@ -127,7 +127,7 @@ } function highlightCodeBlocks(root: HTMLElement) { - for (const el of root.querySelectorAll('code[data-code-token="block"]')) { + for (const el of root.querySelectorAll('code[data-code-token="code_block"]')) { highlightCodeBlockElement(el); } } @@ -151,7 +151,7 @@ } while (node && node !== rootElement) { - if (node instanceof HTMLElement && node.dataset.codeToken === 'block') { + if (node instanceof HTMLElement && node.dataset.codeToken === 'code_block') { const caret = rangeToTextOffset(rootElement, range); if (highlightCodeBlockElement(node)) { @@ -404,7 +404,7 @@ let node: Node | null = container.parentNode; while (node && node !== rootElement) { - if (node instanceof HTMLElement && node.dataset.codeToken === 'block') { + if (node instanceof HTMLElement && node.dataset.codeToken === 'code_block') { const tail = document.createRange(); tail.setStart(container, offset); @@ -462,7 +462,7 @@ const first = rootElement.firstChild; - if (!(first instanceof HTMLElement) || first.dataset.codeToken !== 'block') return false; + if (!(first instanceof HTMLElement) || first.dataset.codeToken !== 'code_block') return false; const range = safeRange(); @@ -507,7 +507,7 @@ const second = first.nextSibling; - if (!(second instanceof HTMLElement) || second.dataset.codeToken !== 'block') return; + if (!(second instanceof HTMLElement) || second.dataset.codeToken !== 'code_block') return; const range = safeRange(); const onHatch = @@ -787,7 +787,7 @@ } -
+
/* pre-wrap is load-bearing: without it Chromium collapses \n in text nodes and converts them to spaces while typing */ - .chat-form-contenteditable { + .chat-form-input-rich { white-space: pre-wrap; } - .chat-form-contenteditable:global([data-empty='true'])::before { + .chat-form-input-rich:global([data-empty='true'])::before { content: attr(data-placeholder); color: var(--muted-foreground); pointer-events: none; } /* Inline code - mirrors markdown-content.css */ - .chat-form-contenteditable :global(code[data-code-token='inline']) { + .chat-form-input-rich :global(code[data-code-token='code_inline']) { background: var(--muted); color: var(--muted-foreground); padding: 0.125rem 0.375rem; @@ -835,7 +835,7 @@ } /* Fenced code block - mirrors .code-block-wrapper in markdown-content.css */ - .chat-form-contenteditable :global(code[data-code-token='block']) { + .chat-form-input-rich :global(code[data-code-token='code_block']) { display: block; margin: 0.25rem 0; padding: 0.75rem 1rem; diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormCommandPicker.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerCommand.svelte similarity index 100% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormCommandPicker.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerCommand.svelte diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormMentionPicker.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte similarity index 100% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormMentionPicker.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickers.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickers.svelte index 7ba97cf9b4a3..dbe03e2e01ac 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickers.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickers.svelte @@ -1,7 +1,7 @@ - -
@@ -191,7 +191,7 @@ icon={Power} tooltip="Load model" class="h-3 w-3 [@media(pointer:coarse)]:text-muted-foreground" - onclick={() => modelsStore.loadModel(option.model)} + onclick={() => modelsStore.status.load(option.model)} stopPropagationOnClick />
diff --git a/tools/ui/src/lib/components/app/models/ModelsSelectorSheet.svelte b/tools/ui/src/lib/components/app/models/ModelsSelectorSheet.svelte index 7228a2e74a79..0d10dd106ccc 100644 --- a/tools/ui/src/lib/components/app/models/ModelsSelectorSheet.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsSelectorSheet.svelte @@ -72,9 +72,9 @@ {@const triggerLoading = !!triggerModel && (triggerStatus === ServerModelStatus.LOADING || - modelsStore.isModelOperationInProgress(triggerModel))} + modelsStore.status.isOperationInProgress(triggerModel))} {@const triggerLoadPercent = triggerLoading - ? Math.round(modelLoadFraction(modelsStore.getLoadProgress(triggerModel)) * 100) + ? Math.round(modelLoadFraction(modelsStore.status.getLoadProgress(triggerModel)) * 100) : 0} {#if ms.isRouter} diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte index c8b2c814cde9..4233039eff44 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte @@ -52,7 +52,7 @@ void modelsStore .fetch() .then(() => modelsStore.fetchRouterModels()) - .then(() => modelsStore.fetchModalitiesForLoadedModels()) + .then(() => modelsStore.props.fetchModalitiesForLoadedModels()) .then(() => modelsStore.ensureFirstModelSelected()); } }); diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte index 30f2b9b2a7d5..d5d93e11d746 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte @@ -23,13 +23,13 @@ let { fields, localConfig, onConfigChange, onThemeChange }: Props = $props(); let currentModelParams = $derived.by(() => { - void modelsStore.propsCacheVersion; + void modelsStore.props.cacheVersion; if (serverStore.isRouterMode) { const currentModelName = modelsStore.selectedModelName; if (currentModelName) { - const currentModelProps = modelsStore.getModelProps(currentModelName); + const currentModelProps = modelsStore.props.getModelProps(currentModelName); return (currentModelProps?.default_generation_settings?.params ?? {}) as Record< string, diff --git a/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte b/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte index 4ea4285322b1..23736ef1c186 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsMcpServers.svelte @@ -121,11 +121,13 @@ {:else} { - const wasEnabled = conversationsStore.isMcpServerEnabledForChat(server.id); + const wasEnabled = conversationsStore.preferences.isMcpServerEnabledForChat( + server.id + ); - await conversationsStore.toggleMcpServerForChat(server.id); + await conversationsStore.preferences.toggleMcpServerForChat(server.id); if (!wasEnabled) { // Promote the connection so tools/prompts/resources become diff --git a/tools/ui/src/lib/constants/attachment-menu.constants.ts b/tools/ui/src/lib/constants/attachment-menu.constants.ts index 62e03bea6e24..07ca17fad151 100644 --- a/tools/ui/src/lib/constants/attachment-menu.constants.ts +++ b/tools/ui/src/lib/constants/attachment-menu.constants.ts @@ -74,7 +74,7 @@ export const ATTACHMENT_PROMPT_ITEMS: AttachmentMenuItem[] = [ enabledWhen: AttachmentItemEnabledWhen.ALWAYS, icon: Zap, id: AttachmentMenuItemId.MCP_PROMPT, - label: 'MCP Prompt', + label: 'MCP Prompts', visibleWhen: AttachmentItemVisibleWhen.HAS_MCP_PROMPTS_SUPPORT } ]; diff --git a/tools/ui/src/lib/constants/cache.constants.ts b/tools/ui/src/lib/constants/cache.constants.ts index b60792d99524..9c6bfadf8abd 100644 --- a/tools/ui/src/lib/constants/cache.constants.ts +++ b/tools/ui/src/lib/constants/cache.constants.ts @@ -32,13 +32,3 @@ export const MCP_RESOURCE_CACHE = { /** TTL for MCP resource cache entries in milliseconds (5 minutes) */ TTL_MS: 5 * 60 * 1000 } as const; - -/** - * Limits for pruning inactive conversation states held in memory. - */ -export const INACTIVE_CONVERSATION = { - /** Maximum age (in ms) for inactive conversation states before cleanup (30 minutes) */ - MAX_AGE_MS: 30 * 60 * 1000, - /** Maximum number of inactive conversation states to keep in memory */ - MAX_STATES: 10 -} as const; diff --git a/tools/ui/src/lib/constants/url.constants.ts b/tools/ui/src/lib/constants/url.constants.ts index 214c8afbac9f..8df442934688 100644 --- a/tools/ui/src/lib/constants/url.constants.ts +++ b/tools/ui/src/lib/constants/url.constants.ts @@ -1,3 +1,5 @@ +import { UrlProtocol } from '$lib/enums'; + const STD = ['com', 'net', 'org', 'gov', 'edu'] as const; const STD_MIL = [...STD, 'mil'] as const; const ccTLD_PREFIXES: Record = { @@ -184,3 +186,7 @@ export const WILDCARD_PUBLIC_SUFFIXES = buildSuffixSet(WILDCARD_BASES); // Matches one or more trailing "/" characters at the end of a URL/path. export const TRAILING_SLASHES_REGEX = /\/+$/; + +// Protocols that apiFetch treats as absolute and passes through untouched. +// Add a protocol here when a caller needs to fetch an absolute URL with it. +export const API_ABSOLUTE_URL_PROTOCOLS = [UrlProtocol.HTTP, UrlProtocol.HTTPS] as const; diff --git a/tools/ui/src/lib/hooks/use-auto-scroll.svelte.ts b/tools/ui/src/lib/hooks/use-auto-scroll.svelte.ts index 6ebce15dad03..d55574efef8c 100644 --- a/tools/ui/src/lib/hooks/use-auto-scroll.svelte.ts +++ b/tools/ui/src/lib/hooks/use-auto-scroll.svelte.ts @@ -14,18 +14,14 @@ export interface AutoScrollOptions { */ export class AutoScrollController { private _autoScrollEnabled = $state(true); - private _userScrolledUp = $state(false); - private _lastScrollTop = $state(0); - private _scrollInterval: ReturnType | undefined; private _container: HTMLElement | undefined; private _disabled: boolean; + private _lastScrollTop = $state(0); private _mutationObserver: MutationObserver | null = null; - private _rafPending = false; private _observerEnabled = false; - constructor(options: AutoScrollOptions = {}) { - this._disabled = options.disabled ?? false; - } - + private _rafPending = false; + private _scrollInterval: ReturnType | undefined; + private _userScrolledUp = $state(false); get autoScrollEnabled(): boolean { return this._autoScrollEnabled; } @@ -34,33 +30,26 @@ export class AutoScrollController { return this._userScrolledUp; } + constructor(options: AutoScrollOptions = {}) { + this._disabled = options.disabled ?? false; + } + /** - * Binds the controller to a scrollable container element. + * Cleans up resources. Call this in onDestroy or when the component unmounts. */ - setContainer(container: HTMLElement | undefined): void { + destroy(): void { + this.stopInterval(); this._doStopObserving(); - this._container = container; - - if (this._observerEnabled && container && !this._disabled) { - this._doStartObserving(); - } } /** - * Updates the disabled state. + * Enables auto-scroll (e.g., when user sends a message). */ - setDisabled(disabled: boolean): void { - if (this._disabled === disabled) return; - - this._disabled = disabled; + enable(): void { + if (this._disabled) return; - if (disabled) { - this._autoScrollEnabled = false; - this.stopInterval(); - this._doStopObserving(); - } else if (this._observerEnabled && this._container && !this._mutationObserver) { - this._doStartObserving(); - } + this._userScrolledUp = false; + this._autoScrollEnabled = true; } /** @@ -85,6 +74,18 @@ export class AutoScrollController { this._lastScrollTop = scrollTop; } + /** + * Resets scroll state when switching conversations. + */ + resetScrollState(): void { + this._userScrolledUp = false; + this._autoScrollEnabled = !this._disabled; + + if (this._container) { + this._lastScrollTop = this._container.scrollTop; + } + } + /** * Scrolls the container to the bottom instantly. */ @@ -95,24 +96,31 @@ export class AutoScrollController { } /** - * Enables auto-scroll (e.g., when user sends a message). + * Binds the controller to a scrollable container element. */ - enable(): void { - if (this._disabled) return; + setContainer(container: HTMLElement | undefined): void { + this._doStopObserving(); + this._container = container; - this._userScrolledUp = false; - this._autoScrollEnabled = true; + if (this._observerEnabled && container && !this._disabled) { + this._doStartObserving(); + } } /** - * Resets scroll state when switching conversations. + * Updates the disabled state. */ - resetScrollState(): void { - this._userScrolledUp = false; - this._autoScrollEnabled = !this._disabled; + setDisabled(disabled: boolean): void { + if (this._disabled === disabled) return; - if (this._container) { - this._lastScrollTop = this._container.scrollTop; + this._disabled = disabled; + + if (disabled) { + this._autoScrollEnabled = false; + this.stopInterval(); + this._doStopObserving(); + } else if (this._observerEnabled && this._container && !this._mutationObserver) { + this._doStartObserving(); } } @@ -127,6 +135,18 @@ export class AutoScrollController { }, AUTO_SCROLL_INTERVAL); } + /** + * Starts a MutationObserver on the container that auto-scrolls to bottom + * on content changes. More responsive than interval-based polling. + */ + startObserving(): void { + this._observerEnabled = true; + + if (this._container && !this._disabled && !this._mutationObserver) { + this._doStartObserving(); + } + } + /** * Stops the auto-scroll interval. */ @@ -137,6 +157,14 @@ export class AutoScrollController { } } + /** + * Stops the MutationObserver. + */ + stopObserving(): void { + this._observerEnabled = false; + this._doStopObserving(); + } + /** * Updates the auto-scroll interval based on streaming state. * Call this in a $effect to automatically manage the interval. @@ -157,34 +185,6 @@ export class AutoScrollController { } } - /** - * Cleans up resources. Call this in onDestroy or when the component unmounts. - */ - destroy(): void { - this.stopInterval(); - this._doStopObserving(); - } - - /** - * Starts a MutationObserver on the container that auto-scrolls to bottom - * on content changes. More responsive than interval-based polling. - */ - startObserving(): void { - this._observerEnabled = true; - - if (this._container && !this._disabled && !this._mutationObserver) { - this._doStartObserving(); - } - } - - /** - * Stops the MutationObserver. - */ - stopObserving(): void { - this._observerEnabled = false; - this._doStopObserving(); - } - private _doStartObserving(): void { if (!this._container || this._mutationObserver) return; diff --git a/tools/ui/src/lib/hooks/use-chat-screen-active-model.svelte.ts b/tools/ui/src/lib/hooks/use-chat-screen-active-model.svelte.ts index ceffdd8a3a31..b5a5d85ce939 100644 --- a/tools/ui/src/lib/hooks/use-chat-screen-active-model.svelte.ts +++ b/tools/ui/src/lib/hooks/use-chat-screen-active-model.svelte.ts @@ -22,10 +22,10 @@ export function useChatScreenActiveModel() { $effect(() => { if (activeModelId) { - const cached = modelsStore.getModelProps(activeModelId); + const cached = modelsStore.props.getModelProps(activeModelId); if (!cached) { - modelsStore.fetchModelProps(activeModelId).then(() => { + modelsStore.props.fetchModelProps(activeModelId).then(() => { modelPropsVersion++; }); } @@ -36,7 +36,7 @@ export function useChatScreenActiveModel() { if (activeModelId) { void modelPropsVersion; - return modelsStore.modelSupportsAudio(activeModelId); + return modelsStore.props.modelSupportsAudio(activeModelId); } return false; @@ -45,7 +45,7 @@ export function useChatScreenActiveModel() { if (activeModelId) { void modelPropsVersion; - return modelsStore.modelSupportsVideo(activeModelId); + return modelsStore.props.modelSupportsVideo(activeModelId); } return false; @@ -54,7 +54,7 @@ export function useChatScreenActiveModel() { if (activeModelId) { void modelPropsVersion; - return modelsStore.modelSupportsVision(activeModelId); + return modelsStore.props.modelSupportsVision(activeModelId); } return false; diff --git a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts index 07d380224d18..c6d55e3935f7 100644 --- a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts +++ b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts @@ -54,10 +54,10 @@ export function useContextGauge(): UseContextGaugeReturn { const modelId = contextStatsStore.activeModelId; if (modelId && contextStatsStore.isActiveModelLoaded) { - const cached = modelsStore.getModelProps(modelId); + const cached = modelsStore.props.getModelProps(modelId); if (!cached) { - void modelsStore.fetchModelProps(modelId); + void modelsStore.props.fetchModelProps(modelId); } } }); @@ -80,9 +80,9 @@ export function useContextGauge(): UseContextGaugeReturn { if (!modelId || contextStatsStore.isActiveModelLoading) return; try { - await modelsStore.loadModel(modelId); + await modelsStore.status.load(modelId); } catch { - // toast already surfaced by modelsStore.loadModel + // toast already surfaced by modelsStore.status.load } } diff --git a/tools/ui/src/lib/hooks/use-models-selector.svelte.ts b/tools/ui/src/lib/hooks/use-models-selector.svelte.ts index d56eeefcd3d2..7d2770a261ad 100644 --- a/tools/ui/src/lib/hooks/use-models-selector.svelte.ts +++ b/tools/ui/src/lib/hooks/use-models-selector.svelte.ts @@ -47,7 +47,7 @@ export interface UseModelsSelectorReturn { export function useModelsSelector(opts: UseModelsSelectorOptions): UseModelsSelectorReturn { const options = $derived( modelsStore.models.filter((option) => { - const modelProps = modelsStore.getModelProps(option.model); + const modelProps = modelsStore.props.getModelProps(option.model); return modelProps?.ui !== false; }) @@ -103,7 +103,7 @@ export function useModelsSelector(opts: UseModelsSelectorOptions): UseModelsSele if (open) { modelsStore.fetchRouterModels().then(() => { - modelsStore.fetchModalitiesForLoadedModels(); + modelsStore.props.fetchModalitiesForLoadedModels(); }); } @@ -143,8 +143,8 @@ export function useModelsSelector(opts: UseModelsSelectorOptions): UseModelsSele if (!onModelChange && isRouter && !modelsStore.isModelLoaded(option.model)) { isLoadingModel = true; - modelsStore - .loadModel(option.model) + modelsStore.status + .load(option.model) .catch((error) => console.error('Failed to load model:', error)) .finally(() => (isLoadingModel = false)); } diff --git a/tools/ui/src/lib/hooks/use-processing-state.svelte.ts b/tools/ui/src/lib/hooks/use-processing-state.svelte.ts index 37e0748bcbe9..8a6f332f3505 100644 --- a/tools/ui/src/lib/hooks/use-processing-state.svelte.ts +++ b/tools/ui/src/lib/hooks/use-processing-state.svelte.ts @@ -43,7 +43,7 @@ export function useProcessingState(): UseProcessingStateReturn { } // Read directly from the reactive state - return chatStore.activeProcessingState; + return chatStore.processing.activeState; }); $effect(() => { diff --git a/tools/ui/src/lib/hooks/use-reasoning-menu.svelte.ts b/tools/ui/src/lib/hooks/use-reasoning-menu.svelte.ts index 2ff67c9392d7..2cb9c906095e 100644 --- a/tools/ui/src/lib/hooks/use-reasoning-menu.svelte.ts +++ b/tools/ui/src/lib/hooks/use-reasoning-menu.svelte.ts @@ -42,19 +42,20 @@ export function useReasoningMenu(): UseReasoningMenuReturn { }); const modelSupportsThinking = $derived.by(() => { void modelsStore.loadedModelIds; - void modelsStore.propsCacheVersion; + void modelsStore.props.cacheVersion; if (serverStore.isRouterMode) { const modelId = modelsStore.selectedModelName || conversationModel; return ( - modelsStore.checkModelSupportsThinking(modelId ?? '') || modelSupportsThinkingFromMessages + modelsStore.props.checkModelSupportsThinking(modelId ?? '') || + modelSupportsThinkingFromMessages ); } - return modelsStore.supportsThinking || modelSupportsThinkingFromMessages; + return modelsStore.props.supportsThinking || modelSupportsThinkingFromMessages; }); - const currentEffort = $derived(conversationsStore.getReasoningEffort()); + const currentEffort = $derived(conversationsStore.preferences.getReasoningEffort()); const thinkingEnabled = $derived( currentEffort !== ReasoningEffort.OFF && currentEffort !== ReasoningEffort.DEFAULT ); @@ -76,7 +77,7 @@ export function useReasoningMenu(): UseReasoningMenuReturn { return modelSupportsThinking; }, select(level: ReasoningEffortLevel): void { - conversationsStore.setReasoningEffort(level.value as ReasoningEffort); + conversationsStore.preferences.setReasoningEffort(level.value as ReasoningEffort); }, get thinkingEnabled() { return thinkingEnabled; diff --git a/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts b/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts index 80b3b85a996c..e9dc0dcab69c 100644 --- a/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts +++ b/tools/ui/src/lib/hooks/use-tools-panel.svelte.ts @@ -35,7 +35,7 @@ export function useToolsPanel(): UseToolsPanelReturn { (g) => g.source !== ToolSource.MCP || !g.serverId || - conversationsStore.isMcpServerEnabledForChat(g.serverId) + conversationsStore.preferences.isMcpServerEnabledForChat(g.serverId) ) ); const totalToolCount = $derived(activeGroups.reduce((n, g) => n + g.tools.length, 0)); @@ -73,7 +73,7 @@ export function useToolsPanel(): UseToolsPanelReturn { return ( group.source === ToolSource.MCP && !!group.serverId && - !conversationsStore.isMcpServerEnabledForChat(group.serverId) + !conversationsStore.preferences.isMcpServerEnabledForChat(group.serverId) ); } diff --git a/tools/ui/src/lib/services/chat.service.ts b/tools/ui/src/lib/services/chat.service.ts index f609b4f4ecad..b008b16db86b 100644 --- a/tools/ui/src/lib/services/chat.service.ts +++ b/tools/ui/src/lib/services/chat.service.ts @@ -1,4 +1,11 @@ -import { settingsStore } from '../stores/settings.svelte'; +/** + * ChatService - Stateless chat completion and streaming API layer + * + * Wraps the /chat/completions and /stream endpoints: request building, SSE + * parsing, streaming callbacks, resume/probe logic and pre-encode KV-cache + * warming. No reactive state; consumed by chatStore and its managers. + */ + import { getAudioInputFormat } from '../utils/audio-format'; import { capImageDataURLSize } from '../utils/cap-img-size'; import { @@ -25,7 +32,8 @@ import { ReasoningFormat, StreamConnectionState } from '$lib/enums'; -import { modelsStore } from '$lib/stores/models.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { DatabaseMessageExtraMcpPrompt, DatabaseMessageExtraMcpResource } from '$lib/types'; import type { ApiChatCompletionToolCall, @@ -53,711 +61,800 @@ function streamStorageKey(conversationId: string): string { } export class ChatService { - /** - * - * - * Title Generation - * - * - */ + // Per-chunk localStorage writes are throttled to at most one per + // conversation per interval (saveStreamStateThrottled). The resume offset + // only needs to be roughly current: on resume the server retransmits from + // a line boundary and the client discards its partial line. Guaranteed + // immediate writes happen at stream start, at resume boundaries and when + // the page goes hidden or away (pagehide/visibilitychange), so a reload + // always finds a usable offset. + private static readonly STREAM_STATE_SAVE_INTERVAL_MS = 500; + + private static streamStateSaveTrackers = new Map< + string, + { lastSavedAt: number; model: string | null; pendingBytes: number | null } + >(); /** - * Sends a streaming chat completion request for generating a chat title. - * Delegates to `sendMessage` for fetch, SSE parsing, and error handling. + * Checks whether all server slots are currently idle (not processing any requests). + * Queries the /slots endpoint (requires --slots flag on the server). + * Returns true if all slots are idle, false if any is processing. + * If the endpoint is unavailable or errors out, returns true (best-effort fallback). * - * @param message - The single message to send (a user message containing the title generation prompt) - * @param model - Optional model name to use (required in ROUTER mode) - * @param signal - Optional AbortSignal to cancel the request - * @returns {Promise} The aggregated title text, or empty string if request failed - * @static + * @param signal - Optional AbortSignal to cancel the request if needed + * @param model - Optional model name to check slots for (required in ROUTER mode) + * @returns {Promise} Promise that resolves to true if all slots are idle, false if any is processing */ - static async generateTitle( - message: ApiChatMessageData, - model?: string | null, - signal?: AbortSignal - ): Promise { - let titleResponse = ''; - + static async areAllSlotsIdle(model?: string | null, signal?: AbortSignal): Promise { try { - await ChatService.sendMessage( - [message], - { - custom: { chat_template_kwargs: { enable_thinking: false } }, - model: model || undefined, - onChunk: (chunk: string) => { - titleResponse += chunk; - }, - stream: true - }, - undefined, - signal - ); + const url = model ? `${API_SLOTS.LIST}?model=${encodeURIComponent(model)}` : API_SLOTS.LIST; + const res = await fetch(url, { signal }); + + if (!res.ok) return true; + + const slots: { is_processing: boolean }[] = await res.json(); + + return slots.every((s) => !s.is_processing); } catch { - return ''; + return true; } - - return titleResponse; } /** - * - * - * Messaging - * - * - */ - - /** - * Sends a chat completion request to the llama-server. - * Supports both streaming and non-streaming responses with comprehensive parameter configuration. - * Automatically converts database messages with attachments to the appropriate API format. - * - * @param messages - Array of chat messages to send to the API (supports both ApiChatMessageData and DatabaseMessage with attachments) - * @param options - Configuration options for the chat completion request. See `SettingsChatServiceOptions` type for details. - * @returns {Promise} that resolves to the complete response string (non-streaming) or void (streaming) - * @throws {Error} if the request fails or is aborted + * Cancels the server-side replay buffer for a conversation, freeing its slot. */ - static async sendMessage( - messages: ApiChatMessageData[] | (DatabaseMessage & { extra?: DatabaseMessageExtra[] })[], - options: SettingsChatServiceOptions = {}, - conversationId?: string, - signal?: AbortSignal - ): Promise { - const { - backend_sampling, - continueFinalMessage, - custom, - // Config options - disableReasoningParsing, - dry_allowed_length, - dry_base, - dry_multiplier, - dry_penalty_last_n, - dynatemp_exponent, - // Sampling parameters - dynatemp_range, - enableThinking, - excludeReasoningFromContext, - frequency_penalty, - max_tokens, - min_p, - onChunk, - onComplete, - onCompletionId, - onConnectionState, - onError, - onModel, - onReasoningChunk, - onTimings, - onToolCallChunk, - presence_penalty, - reasoningEffort, - // Penalty parameters - repeat_last_n, - repeat_penalty, - // Other parameters - samplers, - stream, - // Generation parameters - temperature, - timings_per_token, - // Tools for function calling - tools, - top_k, - top_p, - typ_p, - xtc_probability, - xtc_threshold - } = options; - const normalizedMessages: ApiChatMessageData[] = ( - await Promise.all( - messages.map((msg) => { - if ('id' in msg && 'convId' in msg && 'timestamp' in msg) { - const dbMsg = msg as DatabaseMessage & { extra?: DatabaseMessageExtra[] }; - - return ChatService.convertDbMessageToApiChatMessageData(dbMsg); - } else { - return msg as ApiChatMessageData; - } - }) - ) - ).filter((msg: { role: ChatRole; content: string | ApiChatMessageContentPart[] }) => { - // Filter out empty system messages - if (msg.role === MessageRole.SYSTEM) { - const content = typeof msg.content === 'string' ? msg.content : ''; - - return content.trim().length > 0; - } - - return true; - }); - - // Filter out image attachments if the model doesn't support vision - if (options.model && !modelsStore.modelSupportsVision(options.model)) { - normalizedMessages.forEach((msg) => { - if (Array.isArray(msg.content)) { - msg.content = msg.content.filter((part: ApiChatMessageContentPart) => { - if (part.type === ContentPartType.IMAGE_URL) { - console.info( - `[ChatService] Skipping image attachment in message history (model "${options.model}" does not support vision)` - ); - - return false; - } + static async cancelServerStream(conversationId: string, model?: string | null): Promise { + if (!conversationId) return; - return true; - }); + try { + const id = streamIdentity(conversationId, model); - // If only text remains and it's a single part, simplify to string - if ( - msg.content.length === 1 && - msg.content[0].type === ContentPartType.TEXT && - typeof msg.content[0].text === 'string' - ) { - msg.content = msg.content[0].text; - } - } + await fetch(ChatService.buildStreamUrl(id), { + headers: getAuthHeaders(), + method: 'DELETE' }); + } catch (e) { + console.warn('cancelServerStream failed:', e); } + } - const requestBody: ApiChatCompletionRequest = { - messages: normalizedMessages.map((msg: ApiChatMessageData) => { - const mapped: ApiChatCompletionRequest['messages'][0] = { - content: msg.content, - role: msg.role, - tool_call_id: msg.tool_call_id, - tool_calls: msg.tool_calls - }; + static clearStreamState(conversationId: string): void { + if (!conversationId) return; - // Include reasoning_content from the dedicated field - if (!excludeReasoningFromContext && msg.reasoning_content) { - mapped.reasoning_content = msg.reasoning_content; - } + ChatService.streamStateSaveTrackers.delete(conversationId); - return mapped; - }), - return_progress: stream ? true : undefined, - sse_ping_interval: stream ? 1 : undefined, - stream, - tools: tools && tools.length > 0 ? tools : undefined - }; + try { + localStorage.removeItem(streamStorageKey(conversationId)); + } catch { + // nothing to do + } + } - // Include model in request if provided (required in ROUTER mode) - if (options.model) { - requestBody.model = options.model; + /** + * Converts a database message with attachments to API chat message format. + * Processes various attachment types (images, text files, PDFs) and formats them + * as content parts suitable for the chat completion API. + */ + static async convertDbMessageToApiChatMessageData( + message: DatabaseMessage & { extra?: DatabaseMessageExtra[] } + ): Promise { + // Handle tool result messages (role: 'tool') + if (message.role === MessageRole.TOOL && message.toolCallId) { + return { + content: message.content, + role: MessageRole.TOOL, + tool_call_id: message.toolCallId + }; } - requestBody.reasoning_format = disableReasoningParsing - ? ReasoningFormat.NONE - : ReasoningFormat.AUTO; + // Parse tool calls for assistant messages + let toolCalls: ApiChatCompletionToolCall[] | undefined; - const reasoningBudgetTokens = - enableThinking && reasoningEffort ? (REASONING_EFFORT_TOKENS[reasoningEffort] ?? -1) : -1; + if (message.toolCalls) { + try { + toolCalls = JSON.parse(message.toolCalls); + } catch { + // Ignore parse errors for malformed tool calls + } + } - // an explicit user choice injects the kwarg, otherwise it is omitted so - // the server default applies (--reasoning flag or chat template) - if (enableThinking !== undefined) { - requestBody.chat_template_kwargs = { - ...(requestBody.chat_template_kwargs ?? {}), - enable_thinking: enableThinking + if (!message.extra || message.extra.length === 0) { + const result: ApiChatMessageData = { + content: message.content, + role: message.role as MessageRole }; - } - if (reasoningBudgetTokens >= 0) { - requestBody.thinking_budget_tokens = reasoningBudgetTokens; - } + if (message.reasoningContent) { + result.reasoning_content = message.reasoningContent; + } - // arms the budget sampler so reasoning can be ended at runtime via the control endpoint - requestBody.reasoning_control = true; + if (toolCalls && toolCalls.length > 0) { + result.tool_calls = toolCalls; + } - if (continueFinalMessage) { - requestBody.continue_final_message = true; - requestBody.add_generation_prompt = false; + return result; } - if (temperature !== undefined) requestBody.temperature = temperature; + const contentParts: ApiChatMessageContentPart[] = []; + const textFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraTextFile => + extra.type === AttachmentType.TEXT + ); - if (max_tokens !== undefined) { - // Set max_tokens to -1 (infinite) when explicitly configured as 0 or null - requestBody.max_tokens = max_tokens !== null && max_tokens !== 0 ? max_tokens : -1; + for (const textFile of textFiles) { + contentParts.push({ + text: formatAttachmentText(AttachmentLabel.FILE, textFile.name, textFile.content), + type: ContentPartType.TEXT + }); } - if (dynatemp_range !== undefined) requestBody.dynatemp_range = dynatemp_range; + // Handle legacy 'context' type from the old UI (pasted content) + const legacyContextFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraLegacyContext => + extra.type === AttachmentType.LEGACY_CONTEXT + ); - if (dynatemp_exponent !== undefined) requestBody.dynatemp_exponent = dynatemp_exponent; + for (const legacyContextFile of legacyContextFiles) { + contentParts.push({ + text: formatAttachmentText( + AttachmentLabel.FILE, + legacyContextFile.name, + legacyContextFile.content + ), + type: ContentPartType.TEXT + }); + } - if (top_k !== undefined) requestBody.top_k = top_k; + const imageFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraImageFile => + extra.type === AttachmentType.IMAGE + ); - if (top_p !== undefined) requestBody.top_p = top_p; + for (const image of imageFiles) { + const maxImageResolution = settingsStore.getConfig(SETTINGS_KEYS.MAX_IMAGE_RESOLUTION); + // Caps the resolution and bakes the jpeg exif orientation in one pass, + // untouched images pass through as is + const base64Url = await capImageDataURLSize(image.base64Url, maxImageResolution); - if (min_p !== undefined) requestBody.min_p = min_p; + contentParts.push({ + image_url: { url: base64Url }, + type: ContentPartType.IMAGE_URL + }); + } - if (xtc_probability !== undefined) requestBody.xtc_probability = xtc_probability; + const audioFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraAudioFile => + extra.type === AttachmentType.AUDIO + ); - if (xtc_threshold !== undefined) requestBody.xtc_threshold = xtc_threshold; + for (const audio of audioFiles) { + contentParts.push({ + input_audio: { + data: audio.base64Data, + format: getAudioInputFormat(audio.mimeType) + }, + type: ContentPartType.INPUT_AUDIO + }); + } - if (typ_p !== undefined) requestBody.typ_p = typ_p; + if (message.content) { + contentParts.push({ + text: message.content, + type: ContentPartType.TEXT + }); + } - if (repeat_last_n !== undefined) requestBody.repeat_last_n = repeat_last_n; + const videoFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraVideoFile => + extra.type === AttachmentType.VIDEO + ); - if (repeat_penalty !== undefined) requestBody.repeat_penalty = repeat_penalty; + for (const video of videoFiles) { + contentParts.push({ + input_video: { + data: video.base64Data, + format: video.mimeType.includes('mp4') + ? 'mp4' + : video.mimeType.includes('ogg') + ? 'ogg' + : 'auto' + }, + type: ContentPartType.INPUT_VIDEO + }); + } - if (presence_penalty !== undefined) requestBody.presence_penalty = presence_penalty; + const pdfFiles = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraPdfFile => + extra.type === AttachmentType.PDF + ); - if (frequency_penalty !== undefined) requestBody.frequency_penalty = frequency_penalty; + for (const pdfFile of pdfFiles) { + if (pdfFile.processedAsImages && pdfFile.images) { + for (let i = 0; i < pdfFile.images.length; i++) { + contentParts.push({ + image_url: { url: pdfFile.images[i] }, + type: ContentPartType.IMAGE_URL + }); + } + } else { + contentParts.push({ + text: formatAttachmentText(AttachmentLabel.PDF_FILE, pdfFile.name, pdfFile.content), + type: ContentPartType.TEXT + }); + } + } - if (dry_multiplier !== undefined) requestBody.dry_multiplier = dry_multiplier; + const mcpPrompts = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraMcpPrompt => + extra.type === AttachmentType.MCP_PROMPT + ); - if (dry_base !== undefined) requestBody.dry_base = dry_base; + for (const mcpPrompt of mcpPrompts) { + contentParts.push({ + text: formatAttachmentText( + AttachmentLabel.MCP_PROMPT, + mcpPrompt.name, + mcpPrompt.content, + mcpPrompt.serverName + ), + type: ContentPartType.TEXT + }); + } - if (dry_allowed_length !== undefined) requestBody.dry_allowed_length = dry_allowed_length; + const mcpResources = message.extra.filter( + (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraMcpResource => + extra.type === AttachmentType.MCP_RESOURCE + ); - if (dry_penalty_last_n !== undefined) requestBody.dry_penalty_last_n = dry_penalty_last_n; + for (const mcpResource of mcpResources) { + contentParts.push({ + text: formatAttachmentText( + AttachmentLabel.MCP_RESOURCE, + mcpResource.name, + mcpResource.content, + mcpResource.serverName + ), + type: ContentPartType.TEXT + }); + } - if (samplers !== undefined) { - requestBody.samplers = - typeof samplers === 'string' - ? samplers.split(';').filter((s: string) => s.trim()) - : samplers; + const result: ApiChatMessageData = { + content: contentParts, + role: message.role as MessageRole + }; + + if (message.reasoningContent) { + result.reasoning_content = message.reasoningContent; } - if (backend_sampling !== undefined) requestBody.backend_sampling = backend_sampling; + if (toolCalls && toolCalls.length > 0) { + result.tool_calls = toolCalls; + } - if (timings_per_token !== undefined) requestBody.timings_per_token = timings_per_token; + return result; + } - if (custom) { - try { - const customParams = typeof custom === 'string' ? JSON.parse(custom) : custom; + /** + * Fetch the full replay of a server-side stream from byte 0. Returns the raw Response so the + * caller can pipe it through the SSE parser like a fresh stream. + */ + static async fetchStreamReplay(streamId: string): Promise { + const resp = await fetch(ChatService.buildStreamUrl(streamId, 0), { + headers: getAuthHeaders() + }); - Object.assign(requestBody, customParams); - } catch (error) { - console.warn('Failed to parse custom parameters:', error); - } + if (!resp.ok) { + throw new ApiError(`Stream replay failed with HTTP ${resp.status}`, resp.status); } - try { - const headers: Record = { ...getJsonHeaders() }; + return resp; + } - // tag streaming requests with the conversation id, this single header is the opt in for the - // server side replay buffer and powers discoverActiveStream on tab reopen. with an explicit - // model the ::model suffix keeps the per model session distinct - if (stream && conversationId) { - headers[HEADERS.X_CONVERSATION_ID_HEADER] = streamIdentity(conversationId, options.model); - // persist the pending stream before the fetch: a reload during the model load or - // the prompt processing must still find its way back to the session once it exists - ChatService.saveStreamState(conversationId, 0, options.model ?? null); - } + // write a throttled-but-not-yet-persisted offset immediately; used at + // resume boundaries and on pagehide/visibilitychange so the persisted + // offset is the freshest one when it matters + static flushStreamState(conversationId: string): void { + const tracker = ChatService.streamStateSaveTrackers.get(conversationId); - const response = await fetch(API_CHAT.COMPLETIONS, { - body: JSON.stringify(requestBody), - headers, - method: 'POST', + if (!tracker || tracker.pendingBytes === null) return; + + const { model, pendingBytes } = tracker; + + tracker.lastSavedAt = Date.now(); + tracker.pendingBytes = null; + + ChatService.writeStreamState(conversationId, pendingBytes, model); + } + + /** + * Sends a streaming chat completion request for generating a chat title. + * Delegates to `sendMessage` for fetch, SSE parsing, and error handling. + * + * @param message - The single message to send (a user message containing the title generation prompt) + * @param model - Optional model name to use (required in ROUTER mode) + * @param signal - Optional AbortSignal to cancel the request + * @returns {Promise} The aggregated title text, or empty string if request failed + * @static + */ + static async generateTitle( + message: ApiChatMessageData, + model?: string | null, + signal?: AbortSignal + ): Promise { + let titleResponse = ''; + + try { + await ChatService.sendMessage( + [message], + { + custom: { chat_template_kwargs: { enable_thinking: false } }, + model: model || undefined, + onChunk: (chunk: string) => { + titleResponse += chunk; + }, + stream: true + }, + undefined, signal - }); + ); + } catch { + return ''; + } - if (!response.ok) { - // a rejected request (including one cancelled by a stop during the model load) - // leaves nothing to resume - if (conversationId) { - ChatService.clearStreamState(conversationId); - } + return titleResponse; + } - const error = await ChatService.parseErrorResponse(response); + static getStreamState(conversationId: string): ResumableStreamState | null { + if (!conversationId) return null; - if (onError) { - onError(error); - } + try { + const raw = localStorage.getItem(streamStorageKey(conversationId)); - throw error; - } + if (!raw) return null; - if (stream) { - await ChatService.handleStreamResponse( - response, - onChunk, - onComplete, - onError, - onReasoningChunk, - onToolCallChunk, - onModel, - onCompletionId, - onTimings, - conversationId, - signal, - onConnectionState, - options.model - ); - - return; - } else { - return ChatService.handleNonStreamResponse( - response, - onComplete, - onError, - onToolCallChunk, - onModel - ); - } - } catch (error) { - if (isAbortError(error)) { - console.log('Chat completion request was aborted'); - - return; - } - - let userFriendlyError: Error; - - if (error instanceof Error) { - if (error.name === 'TypeError' && error.message.includes('fetch')) { - userFriendlyError = new Error( - 'Unable to connect to server - please check if the server is running' - ); - userFriendlyError.name = 'NetworkError'; - } else if (error.message.includes('ECONNREFUSED')) { - userFriendlyError = new Error('Connection refused - server may be offline'); - userFriendlyError.name = 'NetworkError'; - } else if (error.message.includes('ETIMEDOUT')) { - userFriendlyError = new Error('Request timed out - the server took too long to respond'); - userFriendlyError.name = 'TimeoutError'; - } else { - userFriendlyError = error; - } - } else { - userFriendlyError = new Error('Unknown error occurred while sending message'); - } - - console.error('Error in sendMessage:', error); + const parsed = JSON.parse(raw) as ResumableStreamState; - if (onError) { - onError(userFriendlyError); - } + if (!parsed || typeof parsed.bytesReceived !== 'number') return null; - throw userFriendlyError; + return parsed; + } catch { + return null; } } /** - * Checks whether all server slots are currently idle (not processing any requests). - * Queries the /slots endpoint (requires --slots flag on the server). - * Returns true if all slots are idle, false if any is processing. - * If the endpoint is unavailable or errors out, returns true (best-effort fallback). - * - * @param signal - Optional AbortSignal to cancel the request if needed - * @param model - Optional model name to check slots for (required in ROUTER mode) - * @returns {Promise} Promise that resolves to true if all slots are idle, false if any is processing + * Handles streaming response from the chat completion API. */ - static async areAllSlotsIdle(model?: string | null, signal?: AbortSignal): Promise { - try { - const url = model ? `${API_SLOTS.LIST}?model=${encodeURIComponent(model)}` : API_SLOTS.LIST; - const res = await fetch(url, { signal }); + static async handleStreamResponse( + response: Response, + onChunk?: (chunk: string) => void, + onComplete?: ( + response: string, + reasoningContent?: string, + timings?: ChatMessageTimings, + toolCalls?: string + ) => void, + onError?: (error: Error) => void, + onReasoningChunk?: (chunk: string) => void, + onToolCallChunk?: (chunk: string) => void, + onModel?: (model: string) => void, + onCompletionId?: (id: string) => void, + onTimings?: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => void, + conversationId?: string, + abortSignal?: AbortSignal, + onConnectionState?: (state: StreamConnectionState) => void, + streamModel?: string | null + ): Promise { + let reader = response.body?.getReader(); - if (!res.ok) return true; + if (!reader) { + throw new Error('No response body'); + } - const slots: { is_processing: boolean }[] = await res.json(); + // bytesParsed is the absolute server side buffer offset of the next byte to parse + // segmentStartOffset is the absolute offset where the current reader started, reset on resume + // segmentBytesRead is wire bytes read by the current reader + let bytesParsed = 0; + let segmentStartOffset = 0; + let segmentBytesRead = 0; + let lastByteAt = Date.now(); + // each resume must produce at least one byte to be retried again + // if a resume returns 200 but yields nothing, we abandon + // since the session has a bounded size, the total number of retries is bounded by construction + let madeProgress = true; - return slots.every((s) => !s.is_processing); - } catch { - return true; + const encoder = new TextEncoder(); + + if (conversationId) { + ChatService.saveStreamState(conversationId, 0, streamModel); } - } - /** - * Ends the current reasoning block of a running completion, targeted by its - * chat completion id (streamed back as `id`). Matching the completion rather - * than a slot index avoids a TOCTOU: a finished completion simply matches - * nothing server side. The model is carried so the router forwards to the - * right child, single model ignores it. Returns true on success. - */ - static async stopReasoning(completionId: string, model?: string | null): Promise { - if (!completionId) { - console.error( - 'stopReasoning: no completion id for the active message, cannot target the running completion' - ); + onConnectionState?.(StreamConnectionState.STREAMING); - return false; - } + let decoder = new TextDecoder(); + let aggregatedContent = ''; + let fullReasoningContent = ''; + let aggregatedToolCalls: ApiChatCompletionToolCall[] = []; + let lastTimings: ChatMessageTimings | undefined; + let streamFinished = false; + let modelEmitted = false; + let idEmitted = false; + let toolCallIndexOffset = 0; + let hasOpenToolCallBatch = false; - const body: Record = { - action: CONTROL_ACTION.END_REASONING, - id: completionId + const finalizeOpenToolCallBatch = () => { + if (!hasOpenToolCallBatch) { + return; + } + + toolCallIndexOffset = aggregatedToolCalls.length; + hasOpenToolCallBatch = false; }; + const processToolCallDelta = (toolCalls?: ApiChatCompletionToolCallDelta[]) => { + if (!toolCalls || toolCalls.length === 0) { + return; + } - if (model) body.model = model; + aggregatedToolCalls = ChatService.mergeToolCallDeltas( + aggregatedToolCalls, + toolCalls, + toolCallIndexOffset + ); - try { - const res = await fetch(API_CHAT.CONTROL, { - body: JSON.stringify(body), - headers: getJsonHeaders(), - method: 'POST' - }); - const data = await res.json().catch(() => null); + if (aggregatedToolCalls.length === 0) { + return; + } - if (!res.ok || data?.success !== true) { - console.error('stopReasoning: control request failed', { - completionId, - response: data, - status: res.status - }); + hasOpenToolCallBatch = true; - return false; - } + const serializedToolCalls = JSON.stringify(aggregatedToolCalls); - return true; - } catch (error) { - console.error('stopReasoning: control request threw', { completionId, error }); + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log('[ChatService] Aggregated tool calls:', serializedToolCalls); + } - return false; - } - } + if (!serializedToolCalls) { + return; + } - /** - * Sends a fire-and-forget request to pre-encode the conversation in the server's KV cache. - * After a response completes, this re-submits the full conversation - * using n_predict=0 and stream=false so the server processes the prompt without generating tokens. - * This warms the cache for the next turn, making it faster. - * - * When excludeReasoningFromContext is true, reasoning content is stripped from the messages - * to match what sendMessage would send on the next turn (avoiding cache misses). - * When false, reasoning_content is preserved so the cached prompt matches the next request. - * - * @param messages - The full conversation including the latest assistant response - * @param model - Optional model name (required in ROUTER mode) - * @param excludeReasoning - Whether to strip reasoning content (should match excludeReasoningFromContext setting) - * @param signal - Optional AbortSignal to cancel the pre-encode request - */ - static async cancelServerStream(conversationId: string, model?: string | null): Promise { - if (!conversationId) return; + if (!abortSignal?.aborted) { + onToolCallChunk?.(serializedToolCalls); + } + }; + const onVisibilityChange = () => { + if (typeof document === 'undefined') return; - try { - const id = streamIdentity(conversationId, model); + if (document.visibilityState === 'hidden') { + // the tab is going to the background and the OS may throttle or + // drop the socket shortly; persist the freshest resume offset now + if (conversationId) ChatService.flushStreamState(conversationId); - await fetch(ChatService.buildStreamUrl(id), { - headers: getAuthHeaders(), - method: 'DELETE' - }); - } catch (e) { - console.warn('cancelServerStream failed:', e); - } - } + return; + } - /** - * Look up server-side stream sessions for the given conversation ids. Ids carry the frozen - * conv::model identity when a model was bound at POST time. - */ - static async lookupStreamSessions(conversationIds: string[]): Promise { - const resp = await fetch(API_STREAM.LOOKUP, { - body: JSON.stringify({ conversation_ids: conversationIds }), - headers: getJsonHeaders(), - method: 'POST' - }); + if (streamFinished) return; - if (!resp.ok) { - throw new ApiError(`Stream lookup failed with HTTP ${resp.status}`, resp.status); - } + if (!conversationId) return; - const body = (await resp.json()) as unknown; + // the bytes have been quiet for too long, the OS likely killed the socket + // kicking the reader unblocks reader.read with done=true so the outer loop can resume + if (Date.now() - lastByteAt > STREAM_VISIBILITY_KICK_MS) { + reader!.cancel().catch(() => {}); + } + }; + const onPageHide = () => { + // a reload or navigation is about to happen; make sure the resume + // offset that getStreamState() will read is not a stale throttled one + if (conversationId) ChatService.flushStreamState(conversationId); + }; - if (!Array.isArray(body)) { - throw new Error('Stream lookup returned a non-array response'); + if (typeof document !== 'undefined') { + document.addEventListener('visibilitychange', onVisibilityChange); + window.addEventListener('pagehide', onPageHide); } - return body as ApiStreamSession[]; - } + try { + let chunk = ''; - /** - * Fetch the full replay of a server-side stream from byte 0. Returns the raw Response so the - * caller can pipe it through the SSE parser like a fresh stream. - */ - static async fetchStreamReplay(streamId: string): Promise { - const resp = await fetch(ChatService.buildStreamUrl(streamId, 0), { - headers: getAuthHeaders() - }); + // outer loop drives the resume cycle, swaps reader on premature end of stream + while (true) { + while (true) { + if (abortSignal?.aborted) break; - if (!resp.ok) { - throw new ApiError(`Stream replay failed with HTTP ${resp.status}`, resp.status); - } + let done: boolean; + let value: Uint8Array | undefined; - return resp; - } + try { + const r = await reader.read(); - /** - * Pick the running session to splice into when discoverActiveStream lists candidates for a - * conversation. Finalized sessions are not candidates: their final content was already written - * to the DB by the original onComplete handler, so attaching to them would replay a buffer that - * may not match what the DB holds. A continue session's buffer holds only the appended deltas, - * not the pre continue prefix, so replaying it as a fresh generation would erase the original. - * - * Among running sessions we tie break on the most recent started_at, which covers the case of - * multiple inferences left running on the same conversation. - */ - static selectActiveStream( - sessions: ApiStreamSession[] | null | undefined - ): ApiStreamSession | null { - if (!Array.isArray(sessions) || sessions.length === 0) { - return null; - } + done = r.done; + value = r.value; + } catch (readErr) { + // reader.read() rejects with TypeError when the underlying connection drops + // instead of just resolving with done=true. treat it like done so the outer + // loop swaps reader via the resume path + if (isAbortError(readErr)) { + throw readErr; + } - const running = sessions.filter((s) => !s.is_done); + console.warn('reader.read() rejected, treating as premature end:', readErr); + done = true; + value = undefined; + } - if (running.length === 0) { - return null; - } + if (done) break; - return running.reduce((best, cur) => (cur.started_at > best.started_at ? cur : best)); - } + if (abortSignal?.aborted) break; - // persist the running byte count and the frozen model for a conversation, a later visit - // resumes the SSE replay at the right offset under the same conv::model identity - static saveStreamState( - conversationId: string, - bytesReceived: number, - model?: string | null - ): void { - if (!conversationId) return; + if (value && value.byteLength > 0) { + segmentBytesRead += value.byteLength; + lastByteAt = Date.now(); - try { - const state: ResumableStreamState = { - bytesReceived, - model: model ?? null, - updatedAt: Date.now() - }; + if (!madeProgress) { + madeProgress = true; + onConnectionState?.(StreamConnectionState.STREAMING); + } + } - localStorage.setItem(streamStorageKey(conversationId), JSON.stringify(state)); - } catch { - // localStorage may be full or disabled, silently ignore - } - } + chunk += decoder.decode(value, { stream: true }); + const lines = chunk.split(SSE_LINE_SEPARATOR); - static getStreamState(conversationId: string): ResumableStreamState | null { - if (!conversationId) return null; + chunk = lines.pop() || ''; - try { - const raw = localStorage.getItem(streamStorageKey(conversationId)); + // the persisted offset must point right after the last fully parsed line, + // the trailing `chunk` is partial bytes still waiting for a newline + if (conversationId) { + const tailBytes = encoder.encode(chunk).byteLength; - if (!raw) return null; + bytesParsed = segmentStartOffset + segmentBytesRead - tailBytes; + ChatService.saveStreamStateThrottled(conversationId, bytesParsed, streamModel); + } - const parsed = JSON.parse(raw) as ResumableStreamState; + for (const line of lines) { + if (abortSignal?.aborted) break; - if (!parsed || typeof parsed.bytesReceived !== 'number') return null; + if (line.startsWith(SSE_DATA_PREFIX)) { + const data = line.slice(SSE_DATA_PREFIX.length).trim(); - return parsed; - } catch { - return null; - } - } + if (data === SSE_DONE_MARKER) { + streamFinished = true; - static clearStreamState(conversationId: string): void { - if (!conversationId) return; + continue; + } - try { - localStorage.removeItem(streamStorageKey(conversationId)); - } catch { - // nothing to do - } - } + try { + const parsed: ApiChatCompletionStreamChunk = JSON.parse(data); + const choice = parsed.choices?.[0]; + const content = choice?.delta?.content; + const reasoningContent = choice?.delta?.reasoning_content; + const toolCalls = choice?.delta?.tool_calls; + const timings = parsed.timings; + const promptProgress = parsed.prompt_progress; + const chunkModel = ChatService.extractModelName(parsed); - /** - * Rebuild the stream identity for a resume. The model persisted at POST time wins, including a - * stored null which means the POST carried no explicit model so the identity stays the bare conv - * id. Only fall back to the caller supplied current model when nothing was persisted. - */ - static resumeStreamIdentity( - conversationId: string, - state: ResumableStreamState | null, - fallbackModel: string | null - ): string { - const model = state && state.model !== undefined ? state.model : fallbackModel; + if (chunkModel && !modelEmitted) { + modelEmitted = true; + onModel?.(chunkModel); + } - return streamIdentity(conversationId, model); - } + if (parsed.id && !idEmitted) { + idEmitted = true; + onCompletionId?.(parsed.id); + } - // build the replay route url for a stream identity, from is the resume byte offset, omitted - // for the cancel route - private static buildStreamUrl(streamId: string, from?: number): string { - const query = `${STREAM_QUERY_PARAMS.CONV_ID}=${encodeURIComponent(streamId)}`; - const offset = from === undefined ? '' : `&${STREAM_QUERY_PARAMS.FROM}=${from}`; + if (promptProgress) { + ChatService.notifyTimings(undefined, promptProgress, onTimings); + } - return `${API_STREAM.BASE}?${query}${offset}`; - } + if (timings) { + ChatService.notifyTimings(timings, promptProgress, onTimings); + lastTimings = timings; + } - /** - * Reconnect to an interrupted stream for this conversation. Returns the fetch Response so the - * existing SSE parser drains it like a fresh stream. The server returns 200 on success, 404 if - * no session exists for the conv_id, and 400 if the offset is below the dropped prefix. - */ - // probe the resume route status without consuming the stream: the SSE route has no HEAD, - // so issue the GET and abort it right after the status line. 0 on network error - static async probeResumeStatus(streamId: string): Promise { - if (!streamId) return 0; + if (content) { + finalizeOpenToolCallBatch(); + aggregatedContent += content; - const ac = new AbortController(); + if (!abortSignal?.aborted) { + onChunk?.(content); + } + } - try { - const resp = await fetch(ChatService.buildStreamUrl(streamId, 0), { - headers: getAuthHeaders(), - signal: ac.signal - }); + if (reasoningContent) { + finalizeOpenToolCallBatch(); + fullReasoningContent += reasoningContent; - ac.abort(); + if (!abortSignal?.aborted) { + onReasoningChunk?.(reasoningContent); + } + } - return resp.status; - } catch { - return 0; - } - } + processToolCallDelta(toolCalls); + } catch (e) { + console.error('Error parsing JSON chunk:', e); + } + } + } - static async resumeStream( - conversationId: string, - signal?: AbortSignal, - model?: string | null - ): Promise { - if (!conversationId) return null; + if (abortSignal?.aborted) break; - const state = ChatService.getStreamState(conversationId); - const from = state?.bytesReceived ?? 0; - const id = streamIdentity(conversationId, model); - const url = ChatService.buildStreamUrl(id, from); + if (streamFinished) break; + } - return await fetch(url, { headers: getAuthHeaders(), method: 'GET', signal }); - } + // inner reader done, decide whether to try a resume + if (abortSignal?.aborted) break; - static async preEncode( - messages: ApiChatMessageData[] | (DatabaseMessage & { extra?: DatabaseMessageExtra[] })[], - model?: string | null, - excludeReasoning?: boolean, - signal?: AbortSignal - ): Promise { - const normalizedMessages: ApiChatMessageData[] = ( - await Promise.all( - messages.map((msg) => { - if ('id' in msg && 'convId' in msg && 'timestamp' in msg) { - return ChatService.convertDbMessageToApiChatMessageData( - msg as DatabaseMessage & { extra?: DatabaseMessageExtra[] } - ); - } + if (streamFinished) break; - return msg as ApiChatMessageData; - }) - ) - ).filter((msg: { role: ChatRole; content: string | ApiChatMessageContentPart[] }) => { - if (msg.role === MessageRole.SYSTEM) { - const content = typeof msg.content === 'string' ? msg.content : ''; + if (!conversationId) break; - return content.trim().length > 0; - } + if (!madeProgress) { + onConnectionState?.(StreamConnectionState.LOST); + onError?.(new Error('Stream resume produced no new bytes, giving up')); - return true; - }); - const requestBody: Record = { - messages: normalizedMessages.map((msg: ApiChatMessageData) => { - const mapped: Record = { - content: excludeReasoning ? ChatService.stripReasoningContent(msg.content) : msg.content, - role: msg.role, - tool_call_id: msg.tool_call_id, - tool_calls: msg.tool_calls - }; + break; + } - if (!excludeReasoning && msg.reasoning_content) { + onConnectionState?.(StreamConnectionState.RESUMING); + madeProgress = false; + + // the server resends starting at bytesParsed, discard any partial line we held, it + // will be retransmitted from a clean line boundary. reuse the frozen model, not the + // live dropdown + // resumeStream reads the offset from localStorage, so persist the + // freshest bytesParsed before asking the server to replay from it + ChatService.flushStreamState(conversationId); + const resumeResp = await ChatService.resumeStream( + conversationId, + abortSignal, + streamModel + ).catch(() => null); + + // an abort landing during the resume request is intentional, not a lost connection + if (abortSignal?.aborted) break; + + if (!resumeResp || resumeResp.status !== 200) { + onConnectionState?.(StreamConnectionState.LOST); + onError?.(new Error('Stream connection lost and could not be resumed')); + + break; + } + + const newReader = resumeResp.body?.getReader(); + + if (!newReader) break; + + try { + reader.releaseLock(); + } catch { + /* ignore */ + } + reader = newReader; + decoder = new TextDecoder(); + chunk = ''; + segmentStartOffset = bytesParsed; + segmentBytesRead = 0; + lastByteAt = Date.now(); + } + + if (abortSignal?.aborted) return; + + if (streamFinished) { + finalizeOpenToolCallBatch(); + + if (conversationId) { + ChatService.clearStreamState(conversationId); + } + + const finalToolCalls = + aggregatedToolCalls.length > 0 ? JSON.stringify(aggregatedToolCalls) : undefined; + + onComplete?.( + aggregatedContent, + fullReasoningContent || undefined, + lastTimings, + finalToolCalls + ); + } + } catch (error) { + const err = error instanceof Error ? error : new Error('Stream error'); + + onError?.(err); + + throw err; + } finally { + if (typeof document !== 'undefined') { + document.removeEventListener('visibilitychange', onVisibilityChange); + window.removeEventListener('pagehide', onPageHide); + } + + try { + reader.releaseLock(); + } catch { + /* ignore */ + } + } + } + + /** + * Look up server-side stream sessions for the given conversation ids. Ids carry the frozen + * conv::model identity when a model was bound at POST time. + */ + static async lookupStreamSessions(conversationIds: string[]): Promise { + const resp = await fetch(API_STREAM.LOOKUP, { + body: JSON.stringify({ conversation_ids: conversationIds }), + headers: getJsonHeaders(), + method: 'POST' + }); + + if (!resp.ok) { + throw new ApiError(`Stream lookup failed with HTTP ${resp.status}`, resp.status); + } + + const body = (await resp.json()) as unknown; + + if (!Array.isArray(body)) { + throw new Error('Stream lookup returned a non-array response'); + } + + return body as ApiStreamSession[]; + } + + /** + * Normalizes an array of messages (database or already-API-shaped) into + * API chat message data, converting DB messages and dropping empty system + * messages. Shared by sendMessage, preEncode and the agentic flow. + */ + static async normalizeMessagesForApi( + messages: ApiChatMessageData[] | (DatabaseMessage & { extra?: DatabaseMessageExtra[] })[] + ): Promise { + return ( + await Promise.all( + messages.map((msg) => { + if ('id' in msg && 'convId' in msg && 'timestamp' in msg) { + return ChatService.convertDbMessageToApiChatMessageData( + msg as DatabaseMessage & { extra?: DatabaseMessageExtra[] } + ); + } + + return msg as ApiChatMessageData; + }) + ) + ).filter((msg: { role: ChatRole; content: string | ApiChatMessageContentPart[] }) => { + // Filter out empty system messages + if (msg.role === MessageRole.SYSTEM) { + const content = typeof msg.content === 'string' ? msg.content : ''; + + return content.trim().length > 0; + } + + return true; + }); + } + + /** + * Fire-and-forget request to pre-encode the conversation in the server's KV cache. + * Re-submits the full conversation with n_predict=0 so the server processes the prompt + * without generating tokens, warming the cache for the next turn. + */ + static async preEncode( + messages: ApiChatMessageData[] | (DatabaseMessage & { extra?: DatabaseMessageExtra[] })[], + model?: string | null, + excludeReasoning?: boolean, + signal?: AbortSignal + ): Promise { + const normalizedMessages: ApiChatMessageData[] = + await ChatService.normalizeMessagesForApi(messages); + const requestBody: Record = { + messages: normalizedMessages.map((msg: ApiChatMessageData) => { + const mapped: Record = { + content: excludeReasoning ? ChatService.stripReasoningContent(msg.content) : msg.content, + role: msg.role, + tool_call_id: msg.tool_call_id, + tool_calls: msg.tool_calls + }; + + if (!excludeReasoning && msg.reasoning_content) { mapped.reasoning_content = msg.reasoning_content; } @@ -785,780 +882,500 @@ export class ChatService { } } - /** - * - * - * Streaming - * - * - */ - - /** - * Handles streaming response from the chat completion API - * @param response - The Response object from the fetch request - * @param onChunk - Optional callback invoked for each content chunk received - * @param onComplete - Optional callback invoked when the stream is complete with full response - * @param onError - Optional callback invoked if an error occurs during streaming - * @param onReasoningChunk - Optional callback invoked for each reasoning content chunk - * @param conversationId - Optional conversation ID for per-conversation state tracking - * @returns {Promise} Promise that resolves when streaming is complete - * @throws {Error} if the stream cannot be read or parsed - */ - static async handleStreamResponse( - response: Response, - onChunk?: (chunk: string) => void, - onComplete?: ( - response: string, - reasoningContent?: string, - timings?: ChatMessageTimings, - toolCalls?: string - ) => void, - onError?: (error: Error) => void, - onReasoningChunk?: (chunk: string) => void, - onToolCallChunk?: (chunk: string) => void, - onModel?: (model: string) => void, - onCompletionId?: (id: string) => void, - onTimings?: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => void, - conversationId?: string, - abortSignal?: AbortSignal, - onConnectionState?: (state: StreamConnectionState) => void, - streamModel?: string | null - ): Promise { - let reader = response.body?.getReader(); + // probe the resume route status without consuming the stream: the SSE route has no HEAD, + // so issue the GET and abort it right after the status line. 0 on network error + static async probeResumeStatus(streamId: string): Promise { + if (!streamId) return 0; - if (!reader) { - throw new Error('No response body'); - } + const ac = new AbortController(); - // bytesParsed is the absolute server side buffer offset of the next byte to parse - // segmentStartOffset is the absolute offset where the current reader started, reset on resume - // segmentBytesRead is wire bytes read by the current reader - let bytesParsed = 0; - let segmentStartOffset = 0; - let segmentBytesRead = 0; - let lastByteAt = Date.now(); - // each resume must produce at least one byte to be retried again - // if a resume returns 200 but yields nothing, we abandon - // since the session has a bounded size, the total number of retries is bounded by construction - let madeProgress = true; + try { + const resp = await fetch(ChatService.buildStreamUrl(streamId, 0), { + headers: getAuthHeaders(), + signal: ac.signal + }); - const encoder = new TextEncoder(); + ac.abort(); - if (conversationId) { - ChatService.saveStreamState(conversationId, 0, streamModel); + return resp.status; + } catch { + return 0; } + } - onConnectionState?.(StreamConnectionState.STREAMING); + static async resumeStream( + conversationId: string, + signal?: AbortSignal, + model?: string | null + ): Promise { + if (!conversationId) return null; - let decoder = new TextDecoder(); - let aggregatedContent = ''; - let fullReasoningContent = ''; - let aggregatedToolCalls: ApiChatCompletionToolCall[] = []; - let lastTimings: ChatMessageTimings | undefined; - let streamFinished = false; - let modelEmitted = false; - let idEmitted = false; - let toolCallIndexOffset = 0; - let hasOpenToolCallBatch = false; + const state = ChatService.getStreamState(conversationId); + const from = state?.bytesReceived ?? 0; + const id = streamIdentity(conversationId, model); + const url = ChatService.buildStreamUrl(id, from); - const finalizeOpenToolCallBatch = () => { - if (!hasOpenToolCallBatch) { - return; - } + return await fetch(url, { headers: getAuthHeaders(), method: 'GET', signal }); + } - toolCallIndexOffset = aggregatedToolCalls.length; - hasOpenToolCallBatch = false; - }; - const processToolCallDelta = (toolCalls?: ApiChatCompletionToolCallDelta[]) => { - if (!toolCalls || toolCalls.length === 0) { - return; - } + /** + * Rebuild the stream identity for a resume. The model persisted at POST time wins, including a + * stored null which means the POST carried no explicit model so the identity stays the bare conv + * id. Only fall back to the caller supplied current model when nothing was persisted. + */ + static resumeStreamIdentity( + conversationId: string, + state: ResumableStreamState | null, + fallbackModel: string | null + ): string { + const model = state && state.model !== undefined ? state.model : fallbackModel; - aggregatedToolCalls = ChatService.mergeToolCallDeltas( - aggregatedToolCalls, - toolCalls, - toolCallIndexOffset - ); - - if (aggregatedToolCalls.length === 0) { - return; - } - - hasOpenToolCallBatch = true; + return streamIdentity(conversationId, model); + } - const serializedToolCalls = JSON.stringify(aggregatedToolCalls); + // persist the running byte count and the frozen model for a conversation, a later visit + // resumes the SSE replay at the right offset under the same conv::model + // identity. Writes immediately; the per-chunk read loop uses the throttled + // variant instead. + static saveStreamState( + conversationId: string, + bytesReceived: number, + model?: string | null + ): void { + if (!conversationId) return; - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log('[ChatService] Aggregated tool calls:', serializedToolCalls); - } + ChatService.writeStreamState(conversationId, bytesReceived, model); + // record the write so a throttled save landing inside the interval + // holds its value pending instead of re-writing + ChatService.streamStateSaveTrackers.set(conversationId, { + lastSavedAt: Date.now(), + model: model ?? null, + pendingBytes: null + }); + } - if (!serializedToolCalls) { - return; - } + // throttled variant for the per-chunk read loop: writes at most once per + // conversation per STREAM_STATE_SAVE_INTERVAL_MS, holding the latest value + // pending until the interval elapses or flushStreamState() forces it out + static saveStreamStateThrottled( + conversationId: string, + bytesReceived: number, + model?: string | null + ): void { + if (!conversationId) return; - if (!abortSignal?.aborted) { - onToolCallChunk?.(serializedToolCalls); - } + const tracker = ChatService.streamStateSaveTrackers.get(conversationId) ?? { + lastSavedAt: 0, + model: null, + pendingBytes: null }; - const onVisibilityChange = () => { - if (typeof document === 'undefined') return; - if (document.visibilityState !== 'visible') return; - - if (streamFinished) return; + tracker.model = model ?? null; - if (!conversationId) return; + if (Date.now() - tracker.lastSavedAt >= ChatService.STREAM_STATE_SAVE_INTERVAL_MS) { + tracker.lastSavedAt = Date.now(); + tracker.pendingBytes = null; + ChatService.writeStreamState(conversationId, bytesReceived, model); + } else { + tracker.pendingBytes = bytesReceived; + } - // the bytes have been quiet for too long, the OS likely killed the socket - // kicking the reader unblocks reader.read with done=true so the outer loop can resume - if (Date.now() - lastByteAt > STREAM_VISIBILITY_KICK_MS) { - reader!.cancel().catch(() => {}); - } - }; + ChatService.streamStateSaveTrackers.set(conversationId, tracker); + } - if (typeof document !== 'undefined') { - document.addEventListener('visibilitychange', onVisibilityChange); + /** + * Pick the running session to splice into when discoverActiveStream lists candidates for a + * conversation. Finalized sessions are not candidates: their final content was already written + * to the DB by the original onComplete handler, so attaching to them would replay a buffer that + * may not match what the DB holds. A continue session's buffer holds only the appended deltas, + * not the pre continue prefix, so replaying it as a fresh generation would erase the original. + * + * Among running sessions we tie break on the most recent started_at, which covers the case of + * multiple inferences left running on the same conversation. + */ + static selectActiveStream( + sessions: ApiStreamSession[] | null | undefined + ): ApiStreamSession | null { + if (!Array.isArray(sessions) || sessions.length === 0) { + return null; } - try { - let chunk = ''; + const running = sessions.filter((s) => !s.is_done); - // outer loop drives the resume cycle, swaps reader on premature end of stream - while (true) { - while (true) { - if (abortSignal?.aborted) break; + if (running.length === 0) { + return null; + } - let done: boolean; - let value: Uint8Array | undefined; + return running.reduce((best, cur) => (cur.started_at > best.started_at ? cur : best)); + } - try { - const r = await reader.read(); + /** + * Sends a chat completion request to the llama-server. + * Supports both streaming and non-streaming responses with comprehensive parameter configuration. + * Automatically converts database messages with attachments to the appropriate API format. + * + * @param messages - Array of chat messages to send to the API (supports both ApiChatMessageData and DatabaseMessage with attachments) + * @param options - Configuration options for the chat completion request. See `SettingsChatServiceOptions` type for details. + * @returns {Promise} that resolves to the complete response string (non-streaming) or void (streaming) + * @throws {Error} if the request fails or is aborted + */ + static async sendMessage( + messages: ApiChatMessageData[] | (DatabaseMessage & { extra?: DatabaseMessageExtra[] })[], + options: SettingsChatServiceOptions = {}, + conversationId?: string, + signal?: AbortSignal + ): Promise { + const { + backend_sampling, + continueFinalMessage, + custom, + // Config options + disableReasoningParsing, + dry_allowed_length, + dry_base, + dry_multiplier, + dry_penalty_last_n, + dynatemp_exponent, + // Sampling parameters + dynatemp_range, + enableThinking, + excludeReasoningFromContext, + frequency_penalty, + max_tokens, + min_p, + onChunk, + onComplete, + onCompletionId, + onConnectionState, + onError, + onModel, + onReasoningChunk, + onTimings, + onToolCallChunk, + presence_penalty, + reasoningEffort, + // Penalty parameters + repeat_last_n, + repeat_penalty, + // Other parameters + samplers, + stream, + // Generation parameters + temperature, + timings_per_token, + // Tools for function calling + tools, + top_k, + top_p, + typ_p, + xtc_probability, + xtc_threshold + } = options; + const normalizedMessages: ApiChatMessageData[] = + await ChatService.normalizeMessagesForApi(messages); - done = r.done; - value = r.value; - } catch (readErr) { - // reader.read() rejects with TypeError when the underlying connection drops - // instead of just resolving with done=true. treat it like done so the outer - // loop swaps reader via the resume path - if (isAbortError(readErr)) { - throw readErr; + // Filter out image attachments if the model doesn't support vision + if (options.model && !modelsStore.props.modelSupportsVision(options.model)) { + normalizedMessages.forEach((msg) => { + if (Array.isArray(msg.content)) { + msg.content = msg.content.filter((part: ApiChatMessageContentPart) => { + if (part.type === ContentPartType.IMAGE_URL) { + console.info( + `[ChatService] Skipping image attachment in message history (model "${options.model}" does not support vision)` + ); + + return false; } - console.warn('reader.read() rejected, treating as premature end:', readErr); - done = true; - value = undefined; + return true; + }); + + // If only text remains and it's a single part, simplify to string + if ( + msg.content.length === 1 && + msg.content[0].type === ContentPartType.TEXT && + typeof msg.content[0].text === 'string' + ) { + msg.content = msg.content[0].text; } + } + }); + } - if (done) break; + const requestBody: ApiChatCompletionRequest = { + messages: normalizedMessages.map((msg: ApiChatMessageData) => { + const mapped: ApiChatCompletionRequest['messages'][0] = { + content: msg.content, + role: msg.role, + tool_call_id: msg.tool_call_id, + tool_calls: msg.tool_calls + }; - if (abortSignal?.aborted) break; + // Include reasoning_content from the dedicated field + if (!excludeReasoningFromContext && msg.reasoning_content) { + mapped.reasoning_content = msg.reasoning_content; + } - if (value && value.byteLength > 0) { - segmentBytesRead += value.byteLength; - lastByteAt = Date.now(); + return mapped; + }), + return_progress: stream ? true : undefined, + sse_ping_interval: stream ? 1 : undefined, + stream, + tools: tools && tools.length > 0 ? tools : undefined + }; - if (!madeProgress) { - madeProgress = true; - onConnectionState?.(StreamConnectionState.STREAMING); - } - } + // Include model in request if provided (required in ROUTER mode) + if (options.model) { + requestBody.model = options.model; + } - chunk += decoder.decode(value, { stream: true }); - const lines = chunk.split(SSE_LINE_SEPARATOR); + requestBody.reasoning_format = disableReasoningParsing + ? ReasoningFormat.NONE + : ReasoningFormat.AUTO; - chunk = lines.pop() || ''; + const reasoningBudgetTokens = + enableThinking && reasoningEffort ? (REASONING_EFFORT_TOKENS[reasoningEffort] ?? -1) : -1; - // the persisted offset must point right after the last fully parsed line, - // the trailing `chunk` is partial bytes still waiting for a newline - if (conversationId) { - const tailBytes = encoder.encode(chunk).byteLength; + // an explicit user choice injects the kwarg, otherwise it is omitted so + // the server default applies (--reasoning flag or chat template) + if (enableThinking !== undefined) { + requestBody.chat_template_kwargs = { + ...(requestBody.chat_template_kwargs ?? {}), + enable_thinking: enableThinking + }; + } - bytesParsed = segmentStartOffset + segmentBytesRead - tailBytes; - ChatService.saveStreamState(conversationId, bytesParsed, streamModel); - } + if (reasoningBudgetTokens >= 0) { + requestBody.thinking_budget_tokens = reasoningBudgetTokens; + } - for (const line of lines) { - if (abortSignal?.aborted) break; + // arms the budget sampler so reasoning can be ended at runtime via the control endpoint + requestBody.reasoning_control = true; - if (line.startsWith(SSE_DATA_PREFIX)) { - const data = line.slice(SSE_DATA_PREFIX.length).trim(); + if (continueFinalMessage) { + requestBody.continue_final_message = true; + requestBody.add_generation_prompt = false; + } - if (data === SSE_DONE_MARKER) { - streamFinished = true; + if (temperature !== undefined) requestBody.temperature = temperature; - continue; - } + if (max_tokens !== undefined) { + // Set max_tokens to -1 (infinite) when explicitly configured as 0 or null + requestBody.max_tokens = max_tokens !== null && max_tokens !== 0 ? max_tokens : -1; + } - try { - const parsed: ApiChatCompletionStreamChunk = JSON.parse(data); - const choice = parsed.choices?.[0]; - const content = choice?.delta?.content; - const reasoningContent = choice?.delta?.reasoning_content; - const toolCalls = choice?.delta?.tool_calls; - const timings = parsed.timings; - const promptProgress = parsed.prompt_progress; - const chunkModel = ChatService.extractModelName(parsed); - - if (chunkModel && !modelEmitted) { - modelEmitted = true; - onModel?.(chunkModel); - } - - if (parsed.id && !idEmitted) { - idEmitted = true; - onCompletionId?.(parsed.id); - } + if (dynatemp_range !== undefined) requestBody.dynatemp_range = dynatemp_range; - if (promptProgress) { - ChatService.notifyTimings(undefined, promptProgress, onTimings); - } + if (dynatemp_exponent !== undefined) requestBody.dynatemp_exponent = dynatemp_exponent; - if (timings) { - ChatService.notifyTimings(timings, promptProgress, onTimings); - lastTimings = timings; - } + if (top_k !== undefined) requestBody.top_k = top_k; - if (content) { - finalizeOpenToolCallBatch(); - aggregatedContent += content; + if (top_p !== undefined) requestBody.top_p = top_p; - if (!abortSignal?.aborted) { - onChunk?.(content); - } - } + if (min_p !== undefined) requestBody.min_p = min_p; - if (reasoningContent) { - finalizeOpenToolCallBatch(); - fullReasoningContent += reasoningContent; + if (xtc_probability !== undefined) requestBody.xtc_probability = xtc_probability; - if (!abortSignal?.aborted) { - onReasoningChunk?.(reasoningContent); - } - } + if (xtc_threshold !== undefined) requestBody.xtc_threshold = xtc_threshold; - processToolCallDelta(toolCalls); - } catch (e) { - console.error('Error parsing JSON chunk:', e); - } - } - } + if (typ_p !== undefined) requestBody.typ_p = typ_p; - if (abortSignal?.aborted) break; + if (repeat_last_n !== undefined) requestBody.repeat_last_n = repeat_last_n; - if (streamFinished) break; - } + if (repeat_penalty !== undefined) requestBody.repeat_penalty = repeat_penalty; - // inner reader done, decide whether to try a resume - if (abortSignal?.aborted) break; + if (presence_penalty !== undefined) requestBody.presence_penalty = presence_penalty; - if (streamFinished) break; + if (frequency_penalty !== undefined) requestBody.frequency_penalty = frequency_penalty; - if (!conversationId) break; + if (dry_multiplier !== undefined) requestBody.dry_multiplier = dry_multiplier; - if (!madeProgress) { - onConnectionState?.(StreamConnectionState.LOST); - onError?.(new Error('Stream resume produced no new bytes, giving up')); + if (dry_base !== undefined) requestBody.dry_base = dry_base; - break; - } + if (dry_allowed_length !== undefined) requestBody.dry_allowed_length = dry_allowed_length; - onConnectionState?.(StreamConnectionState.RESUMING); - madeProgress = false; + if (dry_penalty_last_n !== undefined) requestBody.dry_penalty_last_n = dry_penalty_last_n; - // the server resends starting at bytesParsed, discard any partial line we held, it - // will be retransmitted from a clean line boundary. reuse the frozen model, not the - // live dropdown - const resumeResp = await ChatService.resumeStream( - conversationId, - abortSignal, - streamModel - ).catch(() => null); + if (samplers !== undefined) { + requestBody.samplers = + typeof samplers === 'string' + ? samplers.split(';').filter((s: string) => s.trim()) + : samplers; + } - // an abort landing during the resume request is intentional, not a lost connection - if (abortSignal?.aborted) break; + if (backend_sampling !== undefined) requestBody.backend_sampling = backend_sampling; - if (!resumeResp || resumeResp.status !== 200) { - onConnectionState?.(StreamConnectionState.LOST); - onError?.(new Error('Stream connection lost and could not be resumed')); + if (timings_per_token !== undefined) requestBody.timings_per_token = timings_per_token; - break; - } + if (custom) { + try { + const customParams = typeof custom === 'string' ? JSON.parse(custom) : custom; - const newReader = resumeResp.body?.getReader(); + Object.assign(requestBody, customParams); + } catch (error) { + console.warn('Failed to parse custom parameters:', error); + } + } - if (!newReader) break; + try { + const headers: Record = { ...getJsonHeaders() }; - try { - reader.releaseLock(); - } catch { - /* ignore */ - } - reader = newReader; - decoder = new TextDecoder(); - chunk = ''; - segmentStartOffset = bytesParsed; - segmentBytesRead = 0; - lastByteAt = Date.now(); + // tag streaming requests with the conversation id, this single header is the opt in for the + // server side replay buffer and powers discoverActiveStream on tab reopen. with an explicit + // model the ::model suffix keeps the per model session distinct + if (stream && conversationId) { + headers[HEADERS.X_CONVERSATION_ID_HEADER] = streamIdentity(conversationId, options.model); + // persist the pending stream before the fetch: a reload during the model load or + // the prompt processing must still find its way back to the session once it exists + ChatService.saveStreamState(conversationId, 0, options.model ?? null); } - if (abortSignal?.aborted) return; - - if (streamFinished) { - finalizeOpenToolCallBatch(); + const response = await fetch(API_CHAT.COMPLETIONS, { + body: JSON.stringify(requestBody), + headers, + method: 'POST', + signal + }); + if (!response.ok) { + // a rejected request (including one cancelled by a stop during the model load) + // leaves nothing to resume if (conversationId) { ChatService.clearStreamState(conversationId); } - const finalToolCalls = - aggregatedToolCalls.length > 0 ? JSON.stringify(aggregatedToolCalls) : undefined; + const error = await ChatService.parseErrorResponse(response); - onComplete?.( - aggregatedContent, - fullReasoningContent || undefined, - lastTimings, - finalToolCalls + if (onError) { + onError(error); + } + + throw error; + } + + if (stream) { + await ChatService.handleStreamResponse( + response, + onChunk, + onComplete, + onError, + onReasoningChunk, + onToolCallChunk, + onModel, + onCompletionId, + onTimings, + conversationId, + signal, + onConnectionState, + options.model + ); + + return; + } else { + return ChatService.handleNonStreamResponse( + response, + onComplete, + onError, + onToolCallChunk, + onModel ); } } catch (error) { - const err = error instanceof Error ? error : new Error('Stream error'); + if (isAbortError(error)) { + console.log('Chat completion request was aborted'); - onError?.(err); + return; + } - throw err; - } finally { - if (typeof document !== 'undefined') { - document.removeEventListener('visibilitychange', onVisibilityChange); + let userFriendlyError: Error; + + if (error instanceof Error) { + if (error.name === 'TypeError' && error.message.includes('fetch')) { + userFriendlyError = new Error( + 'Unable to connect to server - please check if the server is running' + ); + userFriendlyError.name = 'NetworkError'; + } else if (error.message.includes('ECONNREFUSED')) { + userFriendlyError = new Error('Connection refused - server may be offline'); + userFriendlyError.name = 'NetworkError'; + } else if (error.message.includes('ETIMEDOUT')) { + userFriendlyError = new Error('Request timed out - the server took too long to respond'); + userFriendlyError.name = 'TimeoutError'; + } else { + userFriendlyError = error; + } + } else { + userFriendlyError = new Error('Unknown error occurred while sending message'); } - try { - reader.releaseLock(); - } catch { - /* ignore */ + console.error('Error in sendMessage:', error); + + if (onError) { + onError(userFriendlyError); } + + throw userFriendlyError; } } /** - * Handles non-streaming response from the chat completion API. - * Parses the JSON response and extracts the generated content. - * - * @param response - The fetch Response object containing the JSON data - * @param onComplete - Optional callback invoked when response is successfully parsed - * @param onError - Optional callback invoked if an error occurs while parsing - * @returns {Promise} Promise that resolves to the generated content string - * @throws {Error} if the response cannot be parsed or is malformed + * Ends the current reasoning block of a running completion, targeted by its + * chat completion id (streamed back as `id`). Matching the completion rather + * than a slot index avoids a TOCTOU: a finished completion simply matches + * nothing server side. The model is carried so the router forwards to the + * right child, single model ignores it. Returns true on success. */ - private static async handleNonStreamResponse( - response: Response, - onComplete?: ( - response: string, - reasoningContent?: string, - timings?: ChatMessageTimings, - toolCalls?: string - ) => void, - onError?: (error: Error) => void, - onToolCallChunk?: (chunk: string) => void, - onModel?: (model: string) => void - ): Promise { - try { - const responseText = await response.text(); - - if (!responseText.trim()) { - const noResponseError = new Error('No response received from server. Please try again.'); + static async stopReasoning(completionId: string, model?: string | null): Promise { + if (!completionId) { + console.error( + 'stopReasoning: no completion id for the active message, cannot target the running completion' + ); - throw noResponseError; - } + return false; + } - const data: ApiChatCompletionResponse = JSON.parse(responseText); - const responseModel = ChatService.extractModelName(data); - - if (responseModel) { - onModel?.(responseModel); - } - - const content = data.choices[0]?.message?.content || ''; - const reasoningContent = data.choices[0]?.message?.reasoning_content; - const toolCalls = data.choices[0]?.message?.tool_calls; - - let serializedToolCalls: string | undefined; - - if (toolCalls && toolCalls.length > 0) { - const mergedToolCalls = ChatService.mergeToolCallDeltas([], toolCalls); - - if (mergedToolCalls.length > 0) { - serializedToolCalls = JSON.stringify(mergedToolCalls); - - if (serializedToolCalls) { - onToolCallChunk?.(serializedToolCalls); - } - } - } - - if (!content.trim() && !serializedToolCalls) { - const noResponseError = new Error('No response received from server. Please try again.'); - - throw noResponseError; - } - - onComplete?.(content, reasoningContent, undefined, serializedToolCalls); - - return content; - } catch (error) { - const err = error instanceof Error ? error : new Error('Parse error'); - - onError?.(err); - - throw err; - } - } - - /** - * Merges tool call deltas into an existing array of tool calls. - * Handles both existing and new tool calls, updating existing ones and adding new ones. - * - * @param existing - The existing array of tool calls to merge into - * @param deltas - The array of tool call deltas to merge - * @param indexOffset - Optional offset to apply to the index of new tool calls - * @returns {ApiChatCompletionToolCall[]} The merged array of tool calls - */ - private static mergeToolCallDeltas( - existing: ApiChatCompletionToolCall[], - deltas: ApiChatCompletionToolCallDelta[], - indexOffset = 0 - ): ApiChatCompletionToolCall[] { - const result = existing.map((call) => ({ - ...call, - function: call.function ? { ...call.function } : undefined - })); - - for (const delta of deltas) { - const index = - typeof delta.index === 'number' && delta.index >= 0 - ? delta.index + indexOffset - : result.length; - - while (result.length <= index) { - result.push({ function: undefined }); - } - - const target = result[index]!; - - if (delta.id) { - target.id = delta.id; - } - - if (delta.type) { - target.type = delta.type; - } - - if (delta.function) { - const fn = target.function ? { ...target.function } : {}; - - if (delta.function.name) { - fn.name = delta.function.name; - } - - if (delta.function.arguments) { - fn.arguments = (fn.arguments ?? '') + delta.function.arguments; - } - - target.function = fn; - } - } - - return result; - } - - /** - * - * - * Conversion - * - * - */ - - /** - * Converts a database message with attachments to API chat message format. - * Processes various attachment types (images, text files, PDFs) and formats them - * as content parts suitable for the chat completion API. - * - * @param message - Database message object with optional extra attachments - * @param message.content - The text content of the message - * @param message.role - The role of the message sender (user, assistant, system) - * @param message.extra - Optional array of message attachments (images, files, etc.) - * @returns {ApiChatMessageData} object formatted for the chat completion API - * @static - */ - static async convertDbMessageToApiChatMessageData( - message: DatabaseMessage & { extra?: DatabaseMessageExtra[] } - ): Promise { - // Handle tool result messages (role: 'tool') - if (message.role === MessageRole.TOOL && message.toolCallId) { - return { - content: message.content, - role: MessageRole.TOOL, - tool_call_id: message.toolCallId - }; - } - - // Parse tool calls for assistant messages - let toolCalls: ApiChatCompletionToolCall[] | undefined; - - if (message.toolCalls) { - try { - toolCalls = JSON.parse(message.toolCalls); - } catch { - // Ignore parse errors for malformed tool calls - } - } - - if (!message.extra || message.extra.length === 0) { - const result: ApiChatMessageData = { - content: message.content, - role: message.role as MessageRole - }; - - if (message.reasoningContent) { - result.reasoning_content = message.reasoningContent; - } - - if (toolCalls && toolCalls.length > 0) { - result.tool_calls = toolCalls; - } - - return result; - } - - const contentParts: ApiChatMessageContentPart[] = []; - const textFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraTextFile => - extra.type === AttachmentType.TEXT - ); - - for (const textFile of textFiles) { - contentParts.push({ - text: formatAttachmentText(AttachmentLabel.FILE, textFile.name, textFile.content), - type: ContentPartType.TEXT - }); - } - - // Handle legacy 'context' type from the old UI (pasted content) - const legacyContextFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraLegacyContext => - extra.type === AttachmentType.LEGACY_CONTEXT - ); - - for (const legacyContextFile of legacyContextFiles) { - contentParts.push({ - text: formatAttachmentText( - AttachmentLabel.FILE, - legacyContextFile.name, - legacyContextFile.content - ), - type: ContentPartType.TEXT - }); - } - - const imageFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraImageFile => - extra.type === AttachmentType.IMAGE - ); - - for (const image of imageFiles) { - const maxImageResolution = settingsStore.getConfig(SETTINGS_KEYS.MAX_IMAGE_RESOLUTION); - // Caps the resolution and bakes the jpeg exif orientation in one pass, - // untouched images pass through as is - const base64Url = await capImageDataURLSize(image.base64Url, maxImageResolution); - - contentParts.push({ - image_url: { url: base64Url }, - type: ContentPartType.IMAGE_URL - }); - } - - const audioFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraAudioFile => - extra.type === AttachmentType.AUDIO - ); - - for (const audio of audioFiles) { - contentParts.push({ - input_audio: { - data: audio.base64Data, - format: getAudioInputFormat(audio.mimeType) - }, - type: ContentPartType.INPUT_AUDIO - }); - } - - if (message.content) { - contentParts.push({ - text: message.content, - type: ContentPartType.TEXT - }); - } - - const videoFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraVideoFile => - extra.type === AttachmentType.VIDEO - ); - - for (const video of videoFiles) { - contentParts.push({ - input_video: { - data: video.base64Data, - format: video.mimeType.includes('mp4') - ? 'mp4' - : video.mimeType.includes('ogg') - ? 'ogg' - : 'auto' - }, - type: ContentPartType.INPUT_VIDEO - }); - } - - const pdfFiles = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraPdfFile => - extra.type === AttachmentType.PDF - ); - - for (const pdfFile of pdfFiles) { - if (pdfFile.processedAsImages && pdfFile.images) { - for (let i = 0; i < pdfFile.images.length; i++) { - contentParts.push({ - image_url: { url: pdfFile.images[i] }, - type: ContentPartType.IMAGE_URL - }); - } - } else { - contentParts.push({ - text: formatAttachmentText(AttachmentLabel.PDF_FILE, pdfFile.name, pdfFile.content), - type: ContentPartType.TEXT - }); - } - } - - const mcpPrompts = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraMcpPrompt => - extra.type === AttachmentType.MCP_PROMPT - ); - - for (const mcpPrompt of mcpPrompts) { - contentParts.push({ - text: formatAttachmentText( - AttachmentLabel.MCP_PROMPT, - mcpPrompt.name, - mcpPrompt.content, - mcpPrompt.serverName - ), - type: ContentPartType.TEXT - }); - } - - const mcpResources = message.extra.filter( - (extra: DatabaseMessageExtra): extra is DatabaseMessageExtraMcpResource => - extra.type === AttachmentType.MCP_RESOURCE - ); - - for (const mcpResource of mcpResources) { - contentParts.push({ - text: formatAttachmentText( - AttachmentLabel.MCP_RESOURCE, - mcpResource.name, - mcpResource.content, - mcpResource.serverName - ), - type: ContentPartType.TEXT - }); - } - - const result: ApiChatMessageData = { - content: contentParts, - role: message.role as MessageRole + const body: Record = { + action: CONTROL_ACTION.END_REASONING, + id: completionId }; - if (message.reasoningContent) { - result.reasoning_content = message.reasoningContent; - } - - if (toolCalls && toolCalls.length > 0) { - result.tool_calls = toolCalls; - } - - return result; - } - - /** - * - * - * Utilities - * - * - */ - - /** - * Strips legacy inline reasoning content tags from message content. - * Handles both plain string content and multipart content arrays. - */ - private static stripReasoningContent( - content: string | ApiChatMessageContentPart[] - ): string | ApiChatMessageContentPart[] { - const stripFromString = (text: string): string => - text.replace(LEGACY_AGENTIC_REGEX.REASONING_BLOCK, '').trim(); - - if (typeof content === 'string') { - return stripFromString(content); - } - - return content.map((part) => { - if (part.type === ContentPartType.TEXT && part.text) { - return { ...part, text: stripFromString(part.text) }; - } - - return part; - }); - } - - /** - * Parses error response and creates appropriate error with context information - * @param response - HTTP response object - * @returns Promise - Parsed error with context info if available - */ - private static async parseErrorResponse( - response: Response - ): Promise { - try { - const errorText = await response.text(); - const errorData: ApiErrorResponse = JSON.parse(errorText); - const message = errorData.error?.message || 'Unknown server error'; - const error = new Error(message) as Error & { - contextInfo?: { n_prompt_tokens: number; n_ctx: number }; - }; - - error.name = response.status === 400 ? 'ServerError' : 'HttpError'; + if (model) body.model = model; - if (errorData.error && 'n_prompt_tokens' in errorData.error && 'n_ctx' in errorData.error) { - error.contextInfo = { - n_ctx: errorData.error.n_ctx, - n_prompt_tokens: errorData.error.n_prompt_tokens - }; - } + try { + const res = await fetch(API_CHAT.CONTROL, { + body: JSON.stringify(body), + headers: getJsonHeaders(), + method: 'POST' + }); + const data = await res.json().catch(() => null); - return error; - } catch { - const fallback = new Error( - `Server error (${response.status}): ${response.statusText}` - ) as Error & { - contextInfo?: { n_prompt_tokens: number; n_ctx: number }; - }; + if (!res.ok || data?.success !== true) { + console.error('stopReasoning: control request failed', { + completionId, + response: data, + status: res.status + }); - fallback.name = 'HttpError'; + return false; + } - return fallback; + return true; + } catch (error) { + console.error('stopReasoning: control request threw', { completionId, error }); + + return false; } } + // build the replay route url for a stream identity, from is the resume byte offset, omitted + // for the cancel route + private static buildStreamUrl(streamId: string, from?: number): string { + const query = `${STREAM_QUERY_PARAMS.CONV_ID}=${encodeURIComponent(streamId)}`; + const offset = from === undefined ? '' : `&${STREAM_QUERY_PARAMS.FROM}=${from}`; + + return `${API_STREAM.BASE}?${query}${offset}`; + } + /** * Extracts model name from Chat Completions API response data. * Handles various response formats including streaming chunks and final responses. @@ -1614,6 +1431,137 @@ export class ChatService { return undefined; } + /** + * Handles non-streaming response from the chat completion API. + * Parses the JSON response and extracts the generated content. + * + * @param response - The fetch Response object containing the JSON data + * @param onComplete - Optional callback invoked when response is successfully parsed + * @param onError - Optional callback invoked if an error occurs while parsing + * @returns {Promise} Promise that resolves to the generated content string + * @throws {Error} if the response cannot be parsed or is malformed + */ + private static async handleNonStreamResponse( + response: Response, + onComplete?: ( + response: string, + reasoningContent?: string, + timings?: ChatMessageTimings, + toolCalls?: string + ) => void, + onError?: (error: Error) => void, + onToolCallChunk?: (chunk: string) => void, + onModel?: (model: string) => void + ): Promise { + try { + const responseText = await response.text(); + + if (!responseText.trim()) { + const noResponseError = new Error('No response received from server. Please try again.'); + + throw noResponseError; + } + + const data: ApiChatCompletionResponse = JSON.parse(responseText); + const responseModel = ChatService.extractModelName(data); + + if (responseModel) { + onModel?.(responseModel); + } + + const content = data.choices[0]?.message?.content || ''; + const reasoningContent = data.choices[0]?.message?.reasoning_content; + const toolCalls = data.choices[0]?.message?.tool_calls; + + let serializedToolCalls: string | undefined; + + if (toolCalls && toolCalls.length > 0) { + const mergedToolCalls = ChatService.mergeToolCallDeltas([], toolCalls); + + if (mergedToolCalls.length > 0) { + serializedToolCalls = JSON.stringify(mergedToolCalls); + + if (serializedToolCalls) { + onToolCallChunk?.(serializedToolCalls); + } + } + } + + if (!content.trim() && !serializedToolCalls) { + const noResponseError = new Error('No response received from server. Please try again.'); + + throw noResponseError; + } + + onComplete?.(content, reasoningContent, undefined, serializedToolCalls); + + return content; + } catch (error) { + const err = error instanceof Error ? error : new Error('Parse error'); + + onError?.(err); + + throw err; + } + } + + /** + * Merges tool call deltas into an existing array of tool calls. + * Handles both existing and new tool calls, updating existing ones and adding new ones. + * + * @param existing - The existing array of tool calls to merge into + * @param deltas - The array of tool call deltas to merge + * @param indexOffset - Optional offset to apply to the index of new tool calls + * @returns {ApiChatCompletionToolCall[]} The merged array of tool calls + */ + private static mergeToolCallDeltas( + existing: ApiChatCompletionToolCall[], + deltas: ApiChatCompletionToolCallDelta[], + indexOffset = 0 + ): ApiChatCompletionToolCall[] { + const result = existing.map((call) => ({ + ...call, + function: call.function ? { ...call.function } : undefined + })); + + for (const delta of deltas) { + const index = + typeof delta.index === 'number' && delta.index >= 0 + ? delta.index + indexOffset + : result.length; + + while (result.length <= index) { + result.push({ function: undefined }); + } + + const target = result[index]!; + + if (delta.id) { + target.id = delta.id; + } + + if (delta.type) { + target.type = delta.type; + } + + if (delta.function) { + const fn = target.function ? { ...target.function } : {}; + + if (delta.function.name) { + fn.name = delta.function.name; + } + + if (delta.function.arguments) { + fn.arguments = (fn.arguments ?? '') + delta.function.arguments; + } + + target.function = fn; + } + } + + return result; + } + /** * Calls the onTimings callback with timing data from streaming response. * @@ -1633,4 +1581,85 @@ export class ChatService { onTimingsCallback(timings, promptProgress); } + + /** + * Parses error response and creates appropriate error with context information + * @param response - HTTP response object + * @returns Promise - Parsed error with context info if available + */ + private static async parseErrorResponse( + response: Response + ): Promise { + try { + const errorText = await response.text(); + const errorData: ApiErrorResponse = JSON.parse(errorText); + const message = errorData.error?.message || 'Unknown server error'; + const error = new Error(message) as Error & { + contextInfo?: { n_prompt_tokens: number; n_ctx: number }; + }; + + error.name = response.status === 400 ? 'ServerError' : 'HttpError'; + + if (errorData.error && 'n_prompt_tokens' in errorData.error && 'n_ctx' in errorData.error) { + error.contextInfo = { + n_ctx: errorData.error.n_ctx, + n_prompt_tokens: errorData.error.n_prompt_tokens + }; + } + + return error; + } catch { + const fallback = new Error( + `Server error (${response.status}): ${response.statusText}` + ) as Error & { + contextInfo?: { n_prompt_tokens: number; n_ctx: number }; + }; + + fallback.name = 'HttpError'; + + return fallback; + } + } + + /** + * Strips legacy inline reasoning content tags from message content. + * Handles both plain string content and multipart content arrays. + */ + private static stripReasoningContent( + content: string | ApiChatMessageContentPart[] + ): string | ApiChatMessageContentPart[] { + const stripFromString = (text: string): string => + text.replace(LEGACY_AGENTIC_REGEX.REASONING_BLOCK, '').trim(); + + if (typeof content === 'string') { + return stripFromString(content); + } + + return content.map((part) => { + if (part.type === ContentPartType.TEXT && part.text) { + return { ...part, text: stripFromString(part.text) }; + } + + return part; + }); + } + + // write the resume state straight to localStorage, bypassing the throttle + private static writeStreamState( + conversationId: string, + bytesReceived: number, + model?: string | null + ): void { + try { + const state: ResumableStreamState = { + bytesReceived, + model: model ?? null, + updatedAt: Date.now() + }; + + localStorage.setItem(streamStorageKey(conversationId), JSON.stringify(state)); + } catch { + // localStorage may be full or disabled, silently ignore + } + } } diff --git a/tools/ui/src/lib/services/conversation-transfer.service.ts b/tools/ui/src/lib/services/conversation-transfer.service.ts index acef580053ce..40a09477a345 100644 --- a/tools/ui/src/lib/services/conversation-transfer.service.ts +++ b/tools/ui/src/lib/services/conversation-transfer.service.ts @@ -17,103 +17,96 @@ import { strFromU8, strToU8, unzipSync, zipSync } from 'fflate'; export class ConversationTransferService { /** - * - * - * JSONL Session Format - * - * + * Triggers a browser download of the provided exported conversation data + * @param data - The exported conversation payload (a single conversation with its messages) + * @param filename - Filename; if omitted, a deterministic name is generated */ + static downloadConversationFile(data: ExportedConversation, filename?: string): void { + const { conv: conversation, messages: msgs } = data; - /** - * Serializes a session (a conversation with its messages) as JSONL. - * The first line is the session header (a `SessionRecordType.SESSION` record - * carrying the conversation properties); each subsequent line is a single message. - * @param data - The exported conversation payload - * @returns The JSONL string (one record per line) - */ - static serializeSessionToJsonl(data: ExportedConversation): string { - const { conv, messages } = data; - const sessionLine = JSON.stringify({ - harness: EXPORT_CONV.HARNESS, - type: SessionRecordType.SESSION, - ...conv - }); - const messageLines = messages.map((message: DatabaseMessage) => { - // `toolCalls` is stored as a JSON string; drop it when empty, otherwise parse it. - const { toolCalls, ...rest } = message; - const normalized = toolCalls ? { ...rest, toolCalls: JSON.parse(toolCalls) } : rest; + if (!conversation) { + console.error('Invalid data: missing conversation'); - return JSON.stringify({ message: normalized, type: SessionRecordType.MESSAGE }); - }); + return; + } - return [sessionLine, ...messageLines].join(NEWLINE); + const downloadFilename = + filename ?? ConversationTransferService.generateConversationFilename(conversation, msgs); + const jsonl = ConversationTransferService.serializeSessionToJsonl(data); + const blob = new Blob([jsonl], { type: MimeTypeText.JSONL }); + + ConversationTransferService.triggerDownload(blob, downloadFilename); } /** - * Parses the JSONL session format produced by {@link serializeSessionToJsonl}. - * A `SessionRecordType.SESSION` line starts a new session; following - * `SessionRecordType.MESSAGE` lines are appended to it. Supports multiple - * sessions in a single file. - * @param text - The JSONL file contents - * @returns The parsed conversations with their messages + * Triggers a browser download of multiple conversations as a `.zip`, one + * `.jsonl` file per conversation. + * @param data - The conversations to export */ - static parseSessionsJsonl(text: string): ExportedConversation[] { - const sessions: ExportedConversation[] = []; - - let current: ExportedConversation | null = null; - - for (const line of text.split(NEWLINE)) { - const trimmed = line.trim(); - - if (!trimmed) continue; - - const record = JSON.parse(trimmed); + static downloadConversationsArchive(data: ExportedConversation[]): void { + if (data.length === 0) { + console.error('Invalid data: no conversations to export'); - if (record.type === SessionRecordType.SESSION) { - // Drop the discriminator and harness marker; the rest is the conversation. - const conv = { ...record }; + return; + } - delete conv.type; - delete conv.harness; - current = { conv: conv as DatabaseConversation, messages: [] }; - sessions.push(current); - } else if (record.type === SessionRecordType.MESSAGE) { - if (!current) { - throw new Error('Invalid JSONL: message record before any session record'); - } + const usedNames = new Set(); + const files: Record = {}; - const message = record.message as DatabaseMessage; + for (const session of data) { + const baseName = ConversationTransferService.generateConversationFilename( + session.conv, + session.messages + ); - // `toolCalls` is parsed to an array on export; the DB stores it as a string. - if (message.toolCalls !== undefined && typeof message.toolCalls !== 'string') { - message.toolCalls = JSON.stringify(message.toolCalls); - } + // Disambiguate any duplicate filenames within the archive. + let entryName = baseName; + let suffix = 1; - current.messages.push(message); + while (usedNames.has(entryName)) { + entryName = baseName.replace( + new RegExp(`${FileExtensionText.JSONL}$`), + `_${suffix++}${FileExtensionText.JSONL}` + ); } - // Ignore unknown record types for forward compatibility. + usedNames.add(entryName); + + files[entryName] = strToU8(ConversationTransferService.serializeSessionToJsonl(session)); } - return sessions; + const archiveName = `${new Date().toISOString().split(EXPORT_CONV.ISO_DATE_TIME_SEPARATOR)[0]}_conversations${FileExtensionText.ZIP}`; + const zipped = zipSync(files); + const blob = new Blob([zipped], { type: MimeTypeApplication.ZIP }); + + ConversationTransferService.triggerDownload(blob, archiveName); } /** - * Reports whether the text is the JSONL session format, whose first non-empty - * line is a `SessionRecordType.SESSION` record. A legacy JSON export starts - * with an array or an object that has no such discriminator. - * @param text - The file contents + * Generates a sanitized filename for a conversation export + * @param conversation - The conversation metadata + * @param msgs - Optional array of messages belonging to the conversation + * @returns The generated filename string */ - private static isSessionsJsonl(text: string): boolean { - const trimmed = text.trimStart(); - const lineEnd = trimmed.indexOf(NEWLINE); - const firstLine = lineEnd === -1 ? trimmed : trimmed.slice(0, lineEnd); + static generateConversationFilename( + conversation: { id?: string; name?: string }, + msgs?: DatabaseMessage[] + ): string { + const conversationName = (conversation.name ?? '').trim().toLowerCase(); + const sanitizedName = conversationName + .replace(EXPORT_CONV.NON_ALPHANUMERIC_REGEX, EXPORT_CONV.NONALNUM_REPLACEMENT) + .replace(EXPORT_CONV.MULTIPLE_UNDERSCORE_REGEX, '_') + .substring(0, EXPORT_CONV.NAME_SUFFIX_MAX_LENGTH); + // If we have messages, use the timestamp of the newest message + const referenceDate = msgs?.length + ? new Date(Math.max(...msgs.map((m) => m.timestamp))) + : new Date(); + const iso = referenceDate.toISOString().slice(0, EXPORT_CONV.ISO_TIMESTAMP_SLICE); + const formattedDate = iso + .replace(EXPORT_CONV.ISO_DATE_TIME_SEPARATOR, EXPORT_CONV.ISO_DATE_TIME_SEPARATOR_REPLACEMENT) + .replaceAll(EXPORT_CONV.ISO_TIME_SEPARATOR, EXPORT_CONV.ISO_TIME_SEPARATOR_REPLACEMENT); + const trimmedConvId = conversation.id?.slice(0, EXPORT_CONV.ID_TRIM_LENGTH) ?? ''; - try { - return JSON.parse(firstLine).type === SessionRecordType.SESSION; - } catch { - // Not a standalone JSON record, so not the JSONL format. - return false; - } + return `${formattedDate}_conv_${trimmedConvId}_${sanitizedName}${FileExtensionText.JSONL}`; } /** @@ -162,104 +155,95 @@ export class ConversationTransferService { } /** - * - * - * Downloads - * - * + * Parses the JSONL session format produced by {@link serializeSessionToJsonl}. + * A `SessionRecordType.SESSION` line starts a new session; following + * `SessionRecordType.MESSAGE` lines are appended to it. Supports multiple + * sessions in a single file. + * @param text - The JSONL file contents + * @returns The parsed conversations with their messages */ + static parseSessionsJsonl(text: string): ExportedConversation[] { + const sessions: ExportedConversation[] = []; - /** - * Generates a sanitized filename for a conversation export - * @param conversation - The conversation metadata - * @param msgs - Optional array of messages belonging to the conversation - * @returns The generated filename string - */ - static generateConversationFilename( - conversation: { id?: string; name?: string }, - msgs?: DatabaseMessage[] - ): string { - const conversationName = (conversation.name ?? '').trim().toLowerCase(); - const sanitizedName = conversationName - .replace(EXPORT_CONV.NON_ALPHANUMERIC_REGEX, EXPORT_CONV.NONALNUM_REPLACEMENT) - .replace(EXPORT_CONV.MULTIPLE_UNDERSCORE_REGEX, '_') - .substring(0, EXPORT_CONV.NAME_SUFFIX_MAX_LENGTH); - // If we have messages, use the timestamp of the newest message - const referenceDate = msgs?.length - ? new Date(Math.max(...msgs.map((m) => m.timestamp))) - : new Date(); - const iso = referenceDate.toISOString().slice(0, EXPORT_CONV.ISO_TIMESTAMP_SLICE); - const formattedDate = iso - .replace(EXPORT_CONV.ISO_DATE_TIME_SEPARATOR, EXPORT_CONV.ISO_DATE_TIME_SEPARATOR_REPLACEMENT) - .replaceAll(EXPORT_CONV.ISO_TIME_SEPARATOR, EXPORT_CONV.ISO_TIME_SEPARATOR_REPLACEMENT); - const trimmedConvId = conversation.id?.slice(0, EXPORT_CONV.ID_TRIM_LENGTH) ?? ''; + let current: ExportedConversation | null = null; - return `${formattedDate}_conv_${trimmedConvId}_${sanitizedName}${FileExtensionText.JSONL}`; - } + for (const line of text.split(NEWLINE)) { + const trimmed = line.trim(); - /** - * Triggers a browser download of the provided exported conversation data - * @param data - The exported conversation payload (a single conversation with its messages) - * @param filename - Filename; if omitted, a deterministic name is generated - */ - static downloadConversationFile(data: ExportedConversation, filename?: string): void { - const { conv: conversation, messages: msgs } = data; + if (!trimmed) continue; - if (!conversation) { - console.error('Invalid data: missing conversation'); + const record = JSON.parse(trimmed); - return; - } + if (record.type === SessionRecordType.SESSION) { + // Drop the discriminator and harness marker; the rest is the conversation. + const conv = { ...record }; - const downloadFilename = - filename ?? ConversationTransferService.generateConversationFilename(conversation, msgs); - const jsonl = ConversationTransferService.serializeSessionToJsonl(data); - const blob = new Blob([jsonl], { type: MimeTypeText.JSONL }); + delete conv.type; + delete conv.harness; + current = { conv: conv as DatabaseConversation, messages: [] }; + sessions.push(current); + } else if (record.type === SessionRecordType.MESSAGE) { + if (!current) { + throw new Error('Invalid JSONL: message record before any session record'); + } - ConversationTransferService.triggerDownload(blob, downloadFilename); - } + const message = record.message as DatabaseMessage; - /** - * Triggers a browser download of multiple conversations as a `.zip`, one - * `.jsonl` file per conversation. - * @param data - The conversations to export - */ - static downloadConversationsArchive(data: ExportedConversation[]): void { - if (data.length === 0) { - console.error('Invalid data: no conversations to export'); + // `toolCalls` is parsed to an array on export; the DB stores it as a string. + if (message.toolCalls !== undefined && typeof message.toolCalls !== 'string') { + message.toolCalls = JSON.stringify(message.toolCalls); + } - return; + current.messages.push(message); + } + // Ignore unknown record types for forward compatibility. } - const usedNames = new Set(); - const files: Record = {}; - - for (const session of data) { - const baseName = ConversationTransferService.generateConversationFilename( - session.conv, - session.messages - ); + return sessions; + } - // Disambiguate any duplicate filenames within the archive. - let entryName = baseName; - let suffix = 1; + /** + * Serializes a session (a conversation with its messages) as JSONL. + * The first line is the session header (a `SessionRecordType.SESSION` record + * carrying the conversation properties); each subsequent line is a single message. + * @param data - The exported conversation payload + * @returns The JSONL string (one record per line) + */ + static serializeSessionToJsonl(data: ExportedConversation): string { + const { conv, messages } = data; + const sessionLine = JSON.stringify({ + harness: EXPORT_CONV.HARNESS, + type: SessionRecordType.SESSION, + ...conv + }); + const messageLines = messages.map((message: DatabaseMessage) => { + // `toolCalls` is stored as a JSON string; drop it when empty, otherwise parse it. + const { toolCalls, ...rest } = message; + const normalized = toolCalls ? { ...rest, toolCalls: JSON.parse(toolCalls) } : rest; - while (usedNames.has(entryName)) { - entryName = baseName.replace( - new RegExp(`${FileExtensionText.JSONL}$`), - `_${suffix++}${FileExtensionText.JSONL}` - ); - } - usedNames.add(entryName); + return JSON.stringify({ message: normalized, type: SessionRecordType.MESSAGE }); + }); - files[entryName] = strToU8(ConversationTransferService.serializeSessionToJsonl(session)); - } + return [sessionLine, ...messageLines].join(NEWLINE); + } - const archiveName = `${new Date().toISOString().split(EXPORT_CONV.ISO_DATE_TIME_SEPARATOR)[0]}_conversations${FileExtensionText.ZIP}`; - const zipped = zipSync(files); - const blob = new Blob([zipped], { type: MimeTypeApplication.ZIP }); + /** + * Reports whether the text is the JSONL session format, whose first non-empty + * line is a `SessionRecordType.SESSION` record. A legacy JSON export starts + * with an array or an object that has no such discriminator. + * @param text - The file contents + */ + private static isSessionsJsonl(text: string): boolean { + const trimmed = text.trimStart(); + const lineEnd = trimmed.indexOf(NEWLINE); + const firstLine = lineEnd === -1 ? trimmed : trimmed.slice(0, lineEnd); - ConversationTransferService.triggerDownload(blob, archiveName); + try { + return JSON.parse(firstLine).type === SessionRecordType.SESSION; + } catch { + // Not a standalone JSON record, so not the JSONL format. + return false; + } } /** diff --git a/tools/ui/src/lib/services/database.service.ts b/tools/ui/src/lib/services/database.service.ts index 89dc58b005dd..a466f8483154 100644 --- a/tools/ui/src/lib/services/database.service.ts +++ b/tools/ui/src/lib/services/database.service.ts @@ -1,3 +1,11 @@ +/** + * DatabaseService - IndexedDB persistence for conversations and messages + * + * Thin Dexie layer over the conversations/messages tables: CRUD, tree + * navigation (descendants, reparenting) and cascading deletes. No reactive + * state; consumed by conversationsStore and the chat flows. + */ + import { IDXDB_STORES, IDXDB_TABLES, STORAGE_APP_NAME } from '$lib/constants'; import { MessageRole } from '$lib/enums'; import type { McpServerOverride } from '$lib/types/database'; @@ -20,12 +28,99 @@ const db = new LlamaUiDatabase(); export class DatabaseService { /** + * Deletes multiple conversations in a single transaction. Each deleted + * conversation has its direct children reparented to the nearest surviving + * ancestor (or promoted to top-level). Children also in `ids` are dropped + * entirely rather than reparented. * + * @param ids - Conversation IDs to delete + */ + static async bulkDeleteConversations(ids: string[]): Promise { + const cleanIds = ids.filter((id): id is string => typeof id === 'string' && id.length > 0); + + if (cleanIds.length === 0) return; + + const idSet = new Set(cleanIds); + + await db.transaction( + 'rw', + [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], + async () => { + // Pre-load each to-delete conversation so the per-id reparent + // walk-up doesn't ping-pong the same ancestry chain. + const prefetched = new Map(); + + let frontier = [...cleanIds]; + + const requested = new Set(frontier); + + while (frontier.length > 0) { + const fetched = await db[IDXDB_TABLES.conversations].bulkGet(frontier); + + frontier = []; + for (let i = 0; i < fetched.length; i++) { + const conv = fetched[i]; + + if (!conv || !conv.id) continue; + + prefetched.set(conv.id, conv); + const ancestor = conv.forkedFromConversationId; + + if (ancestor && !prefetched.has(ancestor) && !requested.has(ancestor)) { + frontier.push(ancestor); + requested.add(ancestor); + } + } + } + + for (const id of cleanIds) { + await this.reparentDirectChildren(id, idSet, prefetched); + } + + await db[IDXDB_TABLES.conversations].bulkDelete(cleanIds); + await db[IDXDB_TABLES.messages].where('convId').anyOf(cleanIds).delete(); + } + ); + } + + /** + * Toggles the pinned status of each conversation in `ids` inside a single + * transaction. Treats `pinned === undefined` as `false`, matching the + * semantics of {@link toggleConversationPin} where `!undefined` evaluates + * to `true`. Returns the resulting pinned state for every id that was + * updated; missing ids are omitted from the map. * - * Conversations - * - * + * @param ids - Conversation IDs to toggle + * @returns Map of id -> new pinned state */ + static async bulkToggleConversationPins(ids: string[]): Promise> { + const cleanIds = ids.filter((id): id is string => typeof id === 'string' && id.length > 0); + const result = new Map(); + + if (cleanIds.length === 0) return result; + + await db.transaction('rw', db[IDXDB_TABLES.conversations], async () => { + const convs = await db[IDXDB_TABLES.conversations].bulkGet(cleanIds); + const updates: DatabaseConversation[] = []; + + for (let i = 0; i < cleanIds.length; i++) { + const conv = convs[i]; + + if (!conv) continue; + + const newPinned = !conv.pinned; + + updates.push({ ...conv, pinned: newPinned }); + result.set(cleanIds[i], newPinned); + } + + if (updates.length === 0) return; + + await db[IDXDB_TABLES.conversations].bulkPut(updates); + }); + + return result; + } /** * Creates a new conversation. @@ -51,14 +146,6 @@ export class DatabaseService { return conversation; } - /** - * - * - * Messages - * - * - */ - /** * Creates a new message branch by adding a message and updating parent/child relationships. * Also updates the conversation's currNode to point to the new message. @@ -96,13 +183,7 @@ export class DatabaseService { // Update parent's children array if parent exists if (parentId !== null) { - const parentMessage = await db[IDXDB_TABLES.messages].get(parentId); - - if (parentMessage) { - await db[IDXDB_TABLES.messages].update(parentId, { - children: [...parentMessage.children, newMessage.id] - }); - } + await this.addChildToParent(parentId, newMessage.id); } await this.updateConversation(message.convId, { @@ -178,9 +259,7 @@ export class DatabaseService { }; await db[IDXDB_TABLES.messages].add(systemMessage); - await db[IDXDB_TABLES.messages].update(parentId, { - children: [...parentMessage.children, systemMessage.id] - }); + await this.addChildToParent(parentId, systemMessage.id); return systemMessage; }); @@ -230,121 +309,6 @@ export class DatabaseService { ); } - /** - * Reparents direct children of `parentId` to the nearest surviving - * ancestor (or promotes them to top-level when the immediate parent was - * top-level). Walking skips any ancestor listed in `excludeIds`, since - * those will be deleted in the same batch — leaving a grandchild pointing - * at an `excludeIds` entry would orphan it. Children whose own id is in - * `excludeIds` are dropped from the updates (the bulk-delete pass will - * remove them). `prefetched` may carry a pre-fetched ancestor map to - * avoid repeat reads inside a bulk transaction. - */ - private static async reparentDirectChildren( - parentId: string, - excludeIds: ReadonlySet = new Set(), - prefetched?: ReadonlyMap - ): Promise { - const conv = prefetched?.get(parentId) ?? (await db[IDXDB_TABLES.conversations].get(parentId)); - - if (!conv) return; - - let newParent = conv.forkedFromConversationId; - - const visited = new Set([parentId]); - - while (newParent && excludeIds.has(newParent)) { - if (visited.has(newParent)) { - newParent = undefined; - - break; - } - - visited.add(newParent); - const next = - prefetched?.get(newParent) ?? (await db[IDXDB_TABLES.conversations].get(newParent)); - - if (!next) { - newParent = undefined; - - break; - } - - newParent = next.forkedFromConversationId; - } - - const directChildren = await db[IDXDB_TABLES.conversations] - .filter((c) => c.forkedFromConversationId === parentId) - .toArray(); - const updates: DatabaseConversation[] = []; - - for (const child of directChildren) { - if (excludeIds.has(child.id)) continue; - - updates.push({ ...child, forkedFromConversationId: newParent }); - } - - if (updates.length === 0) return; - - await db[IDXDB_TABLES.conversations].bulkPut(updates); - } - - /** - * Deletes multiple conversations in a single transaction. Each deleted - * conversation has its direct children reparented to the nearest surviving - * ancestor (or promoted to top-level). Children also in `ids` are dropped - * entirely rather than reparented. - * - * @param ids - Conversation IDs to delete - */ - static async bulkDeleteConversations(ids: string[]): Promise { - const cleanIds = ids.filter((id): id is string => typeof id === 'string' && id.length > 0); - - if (cleanIds.length === 0) return; - - const idSet = new Set(cleanIds); - - await db.transaction( - 'rw', - [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], - async () => { - // Pre-load each to-delete conversation so the per-id reparent - // walk-up doesn't ping-pong the same ancestry chain. - const prefetched = new Map(); - - let frontier = [...cleanIds]; - - const requested = new Set(frontier); - - while (frontier.length > 0) { - const fetched = await db[IDXDB_TABLES.conversations].bulkGet(frontier); - - frontier = []; - for (let i = 0; i < fetched.length; i++) { - const conv = fetched[i]; - - if (!conv || !conv.id) continue; - - prefetched.set(conv.id, conv); - const ancestor = conv.forkedFromConversationId; - - if (ancestor && !prefetched.has(ancestor) && !requested.has(ancestor)) { - frontier.push(ancestor); - requested.add(ancestor); - } - } - } - - for (const id of cleanIds) { - await this.reparentDirectChildren(id, idSet, prefetched); - } - - await db[IDXDB_TABLES.conversations].bulkDelete(cleanIds); - await db[IDXDB_TABLES.messages].where('convId').anyOf(cleanIds).delete(); - } - ); - } - /** * Deletes a message and removes it from its parent's children array. * @@ -356,17 +320,8 @@ export class DatabaseService { if (!message) return; - // Remove this message from its parent's children array - if (message.parent) { - const parent = await db[IDXDB_TABLES.messages].get(message.parent); - - if (parent) { - parent.children = parent.children.filter((childId: string) => childId !== messageId); - await db[IDXDB_TABLES.messages].put(parent); - } - } + await this.removeChildFromParent(messageId); - // Delete the message await db[IDXDB_TABLES.messages].delete(messageId); }); } @@ -389,20 +344,10 @@ export class DatabaseService { .where('convId') .equals(conversationId) .toArray(); - // Find all descendant messages const descendants = findDescendantMessages(allMessages, messageId); const allToDelete = [messageId, ...descendants]; - // Get the message to delete for parent cleanup - const message = await db[IDXDB_TABLES.messages].get(messageId); - if (message && message.parent) { - const parent = await db[IDXDB_TABLES.messages].get(message.parent); - - if (parent) { - parent.children = parent.children.filter((childId: string) => childId !== messageId); - await db[IDXDB_TABLES.messages].put(parent); - } - } + await this.removeChildFromParent(messageId); // Delete all messages in the branch await db[IDXDB_TABLES.messages].bulkDelete(allToDelete); @@ -412,23 +357,108 @@ export class DatabaseService { } /** - * Gets all conversations, sorted by last modified time (newest first). - * - * @returns Array of conversations - */ - static async getAllConversations(): Promise { - return await db[IDXDB_TABLES.conversations].orderBy('lastModified').reverse().toArray(); - } - - /** - * Gets a conversation by ID. + * Forks a conversation at a specific message, creating a new conversation + * containing all messages from the root up to (and including) the target message. * - * @param id - Conversation ID - * @returns The conversation if found, otherwise undefined + * @param sourceConvId - The source conversation ID + * @param atMessageId - The message ID to fork at (the new conversation ends here) + * @param options - Fork options (name and whether to include attachments) + * @returns The newly created conversation */ - static async getConversation(id: string): Promise { - return await db[IDXDB_TABLES.conversations].get(id); - } + static async forkConversation( + sourceConvId: string, + atMessageId: string, + options: { name: string; includeAttachments: boolean } + ): Promise { + return await db.transaction( + 'rw', + [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], + async () => { + const sourceConv = await db[IDXDB_TABLES.conversations].get(sourceConvId); + + if (!sourceConv) { + throw new Error(`Source conversation ${sourceConvId} not found`); + } + + const allMessages = await db[IDXDB_TABLES.messages] + .where('convId') + .equals(sourceConvId) + .toArray(); + const pathMessages = filterByLeafNodeId( + allMessages, + atMessageId, + true + ) as DatabaseMessage[]; + + if (pathMessages.length === 0) { + throw new Error(`Could not resolve message path to ${atMessageId}`); + } + + const idMap = new Map(); + + for (const msg of pathMessages) { + idMap.set(msg.id, uuid()); + } + + const newConvId = uuid(); + const clonedMessages: DatabaseMessage[] = pathMessages.map((msg) => { + const newId = idMap.get(msg.id)!; + const newParent = msg.parent ? (idMap.get(msg.parent) ?? null) : null; + const newChildren = msg.children + .filter((childId: string) => idMap.has(childId)) + .map((childId: string) => idMap.get(childId)!); + + return { + ...msg, + children: newChildren, + convId: newConvId, + extra: options.includeAttachments ? msg.extra : undefined, + id: newId, + parent: newParent + }; + }); + const lastClonedMessage = clonedMessages[clonedMessages.length - 1]; + const newConv: DatabaseConversation = { + currNode: lastClonedMessage.id, + cwd: sourceConv.cwd, + forkedFromConversationId: sourceConvId, + id: newConvId, + lastModified: Date.now(), + mcpServerOverrides: sourceConv.mcpServerOverrides + ? sourceConv.mcpServerOverrides.map((o: McpServerOverride) => ({ + enabled: o.enabled, + serverId: o.serverId + })) + : undefined, + name: options.name + }; + + await db[IDXDB_TABLES.conversations].add(newConv); + await db[IDXDB_TABLES.messages].bulkAdd(clonedMessages); + + return newConv; + } + ); + } + + /** + * Gets all conversations, sorted by last modified time (newest first). + * + * @returns Array of conversations + */ + static async getAllConversations(): Promise { + return await db[IDXDB_TABLES.conversations].orderBy('lastModified').reverse().toArray(); + } + + /** + * Gets a conversation by ID. + * + * @param id - Conversation ID + * @returns The conversation if found, otherwise undefined + */ + static async getConversation(id: string): Promise { + return await db[IDXDB_TABLES.conversations].get(id); + } /** * Gets all messages in a conversation, sorted by timestamp (oldest first). @@ -484,27 +514,44 @@ export class DatabaseService { } /** - * Updates a conversation. `lastModified` is never stamped implicitly; - * pass it in `updates` to bump the conversation in recency ordering. + * Imports multiple conversations and their messages. + * Skips conversations that already exist. * - * @param id - Conversation ID - * @param updates - Partial updates to apply - * @returns Promise that resolves when the conversation is updated + * @param data - Array of { conv, messages } objects + * @returns The conversations written to the database and the ones skipped */ - static async updateConversation( - id: string, - updates: Partial> - ): Promise { - await db[IDXDB_TABLES.conversations].update(id, updates); - } + static async importConversations( + data: { conv: DatabaseConversation; messages: DatabaseMessage[] }[] + ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { + const imported: DatabaseConversation[] = []; + const skipped: DatabaseConversation[] = []; - /** - * - * - * Navigation - * - * - */ + return await db.transaction( + 'rw', + [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], + async () => { + for (const item of data) { + const { conv, messages } = item; + const existing = await db[IDXDB_TABLES.conversations].get(conv.id); + + if (existing) { + skipped.push(conv); + + continue; + } + + await db[IDXDB_TABLES.conversations].add(conv); + for (const msg of messages) { + await db[IDXDB_TABLES.messages].put(msg); + } + + imported.push(conv); + } + + return { imported, skipped }; + } + ); + } /** * Toggles the pinned status of a conversation. @@ -527,42 +574,18 @@ export class DatabaseService { } /** - * Toggles the pinned status of each conversation in `ids` inside a single - * transaction. Treats `pinned === undefined` as `false`, matching the - * semantics of {@link toggleConversationPin} where `!undefined` evaluates - * to `true`. Returns the resulting pinned state for every id that was - * updated; missing ids are omitted from the map. + * Updates a conversation. `lastModified` is never stamped implicitly; + * pass it in `updates` to bump the conversation in recency ordering. * - * @param ids - Conversation IDs to toggle - * @returns Map of id -> new pinned state + * @param id - Conversation ID + * @param updates - Partial updates to apply + * @returns Promise that resolves when the conversation is updated */ - static async bulkToggleConversationPins(ids: string[]): Promise> { - const cleanIds = ids.filter((id): id is string => typeof id === 'string' && id.length > 0); - const result = new Map(); - - if (cleanIds.length === 0) return result; - - await db.transaction('rw', db[IDXDB_TABLES.conversations], async () => { - const convs = await db[IDXDB_TABLES.conversations].bulkGet(cleanIds); - const updates: DatabaseConversation[] = []; - - for (let i = 0; i < cleanIds.length; i++) { - const conv = convs[i]; - - if (!conv) continue; - - const newPinned = !conv.pinned; - - updates.push({ ...conv, pinned: newPinned }); - result.set(cleanIds[i], newPinned); - } - - if (updates.length === 0) return; - - await db[IDXDB_TABLES.conversations].bulkPut(updates); - }); - - return result; + static async updateConversation( + id: string, + updates: Partial> + ): Promise { + await db[IDXDB_TABLES.conversations].update(id, updates); } /** @@ -593,146 +616,90 @@ export class DatabaseService { } /** - * - * - * Import - * - * + * Appends a child id to a parent message's children array. */ + private static async addChildToParent(parentId: string, childId: string): Promise { + const parent = await db[IDXDB_TABLES.messages].get(parentId); + + if (!parent) return; + + await db[IDXDB_TABLES.messages].update(parentId, { + children: [...parent.children, childId] + }); + } /** - * Imports multiple conversations and their messages. - * Skips conversations that already exist. - * - * @param data - Array of { conv, messages } objects - * @returns The conversations written to the database and the ones skipped + * Removes a child id from its parent message's children array. */ - static async importConversations( - data: { conv: DatabaseConversation; messages: DatabaseMessage[] }[] - ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { - const imported: DatabaseConversation[] = []; - const skipped: DatabaseConversation[] = []; + private static async removeChildFromParent(messageId: string): Promise { + const message = await db[IDXDB_TABLES.messages].get(messageId); - return await db.transaction( - 'rw', - [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], - async () => { - for (const item of data) { - const { conv, messages } = item; - const existing = await db[IDXDB_TABLES.conversations].get(conv.id); + if (!message?.parent) return; - if (existing) { - skipped.push(conv); - - continue; - } + const parent = await db[IDXDB_TABLES.messages].get(message.parent); - await db[IDXDB_TABLES.conversations].add(conv); - for (const msg of messages) { - await db[IDXDB_TABLES.messages].put(msg); - } + if (!parent) return; - imported.push(conv); - } - - return { imported, skipped }; - } - ); + parent.children = parent.children.filter((childId: string) => childId !== messageId); + await db[IDXDB_TABLES.messages].put(parent); } /** - * - * - * Forking - * - * + * Reparents direct children of `parentId` to the nearest surviving + * ancestor (or promotes them to top-level when the immediate parent was + * top-level). Walking skips any ancestor listed in `excludeIds`, since + * those will be deleted in the same batch — leaving a grandchild pointing + * at an `excludeIds` entry would orphan it. Children whose own id is in + * `excludeIds` are dropped from the updates (the bulk-delete pass will + * remove them). `prefetched` may carry a pre-fetched ancestor map to + * avoid repeat reads inside a bulk transaction. */ + private static async reparentDirectChildren( + parentId: string, + excludeIds: ReadonlySet = new Set(), + prefetched?: ReadonlyMap + ): Promise { + const conv = prefetched?.get(parentId) ?? (await db[IDXDB_TABLES.conversations].get(parentId)); - /** - * Forks a conversation at a specific message, creating a new conversation - * containing all messages from the root up to (and including) the target message. - * - * @param sourceConvId - The source conversation ID - * @param atMessageId - The message ID to fork at (the new conversation ends here) - * @param options - Fork options (name and whether to include attachments) - * @returns The newly created conversation - */ - static async forkConversation( - sourceConvId: string, - atMessageId: string, - options: { name: string; includeAttachments: boolean } - ): Promise { - return await db.transaction( - 'rw', - [db[IDXDB_TABLES.conversations], db[IDXDB_TABLES.messages]], - async () => { - const sourceConv = await db[IDXDB_TABLES.conversations].get(sourceConvId); + if (!conv) return; - if (!sourceConv) { - throw new Error(`Source conversation ${sourceConvId} not found`); - } + let newParent = conv.forkedFromConversationId; - const allMessages = await db[IDXDB_TABLES.messages] - .where('convId') - .equals(sourceConvId) - .toArray(); - const pathMessages = filterByLeafNodeId( - allMessages, - atMessageId, - true - ) as DatabaseMessage[]; + const visited = new Set([parentId]); - if (pathMessages.length === 0) { - throw new Error(`Could not resolve message path to ${atMessageId}`); - } + while (newParent && excludeIds.has(newParent)) { + if (visited.has(newParent)) { + newParent = undefined; - const idMap = new Map(); + break; + } - for (const msg of pathMessages) { - idMap.set(msg.id, uuid()); - } + visited.add(newParent); + const next = + prefetched?.get(newParent) ?? (await db[IDXDB_TABLES.conversations].get(newParent)); - const newConvId = uuid(); - const clonedMessages: DatabaseMessage[] = pathMessages.map((msg) => { - const newId = idMap.get(msg.id)!; - const newParent = msg.parent ? (idMap.get(msg.parent) ?? null) : null; - const newChildren = msg.children - .filter((childId: string) => idMap.has(childId)) - .map((childId: string) => idMap.get(childId)!); + if (!next) { + newParent = undefined; - return { - ...msg, - children: newChildren, - convId: newConvId, - extra: options.includeAttachments ? msg.extra : undefined, - id: newId, - parent: newParent - }; - }); - const lastClonedMessage = clonedMessages[clonedMessages.length - 1]; - const newConv: DatabaseConversation = { - currNode: lastClonedMessage.id, - cwd: sourceConv.cwd, - forkedFromConversationId: sourceConvId, - id: newConvId, - lastModified: Date.now(), - mcpServerOverrides: sourceConv.mcpServerOverrides - ? sourceConv.mcpServerOverrides.map((o: McpServerOverride) => ({ - enabled: o.enabled, - serverId: o.serverId - })) - : undefined, - name: options.name - }; + break; + } - await db[IDXDB_TABLES.conversations].add(newConv); + newParent = next.forkedFromConversationId; + } - for (const msg of clonedMessages) { - await db[IDXDB_TABLES.messages].add(msg); - } + const directChildren = await db[IDXDB_TABLES.conversations] + .filter((c) => c.forkedFromConversationId === parentId) + .toArray(); + const updates: DatabaseConversation[] = []; - return newConv; - } - ); + for (const child of directChildren) { + if (excludeIds.has(child.id)) continue; + + updates.push({ ...child, forkedFromConversationId: newParent }); + } + + if (updates.length === 0) return; + + await db[IDXDB_TABLES.conversations].bulkPut(updates); } } diff --git a/tools/ui/src/lib/services/index.ts b/tools/ui/src/lib/services/index.ts index fe739a3bc257..7ae9e23d48a6 100644 --- a/tools/ui/src/lib/services/index.ts +++ b/tools/ui/src/lib/services/index.ts @@ -53,9 +53,9 @@ * - Reasoning content stripping from prompt history to avoid KV cache pollution * - Error translation (network, timeout, server errors → user-friendly messages) * - * @see chatStore in stores/chat.svelte.ts — primary consumer for chat state management - * @see agenticStore in stores/agentic.svelte.ts — uses ChatService for agentic loop streaming - * @see conversationsStore in stores/conversations.svelte.ts — provides message context + * @see chatStore in stores/chat/index.svelte.ts — primary consumer for chat state management + * @see agenticStore in stores/agentic/index.svelte.ts — uses ChatService for agentic loop streaming + * @see conversationsStore in stores/conversations/index.svelte.ts — provides message context */ export { ChatService } from './chat.service'; @@ -98,8 +98,8 @@ export { ChatService } from './chat.service'; * enabling conversation branching and alternative response paths. The conversation's * `currNode` tracks the currently active branch endpoint. * - * @see conversationsStore in stores/conversations.svelte.ts — reactive layer on top of DatabaseService - * @see chatStore in stores/chat.svelte.ts — uses DatabaseService directly for message CRUD during streaming + * @see conversationsStore in stores/conversations/index.svelte.ts — reactive layer on top of DatabaseService + * @see chatStore in stores/chat/index.svelte.ts — uses DatabaseService directly for message CRUD during streaming */ export { DatabaseService } from './database.service'; @@ -143,7 +143,7 @@ export { ConversationTransferService } from './conversation-transfer.service'; * - `POST /models/load` — Load a model (ROUTER mode only) * - `POST /models/unload` — Unload a model (ROUTER mode only) * - * @see modelsStore in stores/models.svelte.ts — primary consumer for reactive model state + * @see modelsStore in stores/models/index.svelte.ts — primary consumer for reactive model state */ export { ModelsService } from './models.service'; @@ -174,8 +174,8 @@ export { ModelsService } from './models.service'; * - `&autoload=false` → Prevents model auto-loading when querying props * * @see serverStore in stores/server.svelte.ts — consumes global server props - * @see modelsStore in stores/models.svelte.ts — consumes per-model props for modalities - * @see settingsStore in stores/settings.svelte.ts — syncs default generation params from props + * @see modelsStore in stores/models/index.svelte.ts — consumes per-model props for modalities + * @see settingsStore in stores/settings/index.svelte.ts — syncs default generation params from props */ export { PropsService } from './props.service'; @@ -217,7 +217,7 @@ export { PropsService } from './props.service'; * - `ParameterSyncService` class — static methods for sync logic * - `SYNCABLE_PARAMETERS` — mapping of UI setting keys to server parameter keys * - * @see settingsStore in stores/settings.svelte.ts — primary consumer for settings sync + * @see settingsStore in stores/settings/index.svelte.ts — primary consumer for settings sync * @see SettingsChatParameterSourceIndicator — displays parameter source badges in UI */ export { ParameterSyncService } from './parameter-sync.service'; @@ -241,7 +241,7 @@ export { ParameterSyncService } from './parameter-sync.service'; * - Manages connection lifecycle, health checks, reconnection * - Handles tool name conflict resolution and server coordination * - * - **mcpResourceStore**: Reactive resource state + * - **mcpResourceStore** (composed as mcpStore.resources): Reactive resource state * - Receives resource data fetched via MCPService * - Manages resource caching, subscriptions, and attachments * @@ -263,9 +263,9 @@ export { ParameterSyncService } from './parameter-sync.service'; * 2. **StreamableHTTP** — modern HTTP-based, supports CORS proxy * 3. **SSE** — legacy fallback, supports CORS proxy * - * @see mcpStore in stores/mcp.svelte.ts — reactive business logic facade on top of MCPService - * @see mcpResourceStore in stores/mcp-resources.svelte.ts — reactive resource state management - * @see agenticStore in stores/agentic.svelte.ts — uses MCPService (via mcpStore) for tool execution + * @see mcpStore in stores/mcp/index.svelte.ts — reactive business logic facade on top of MCPService + * @see mcpStore.resources in stores/mcp/resources.svelte.ts — reactive resource state management + * @see agenticStore in stores/agentic/index.svelte.ts — uses MCPService (via mcpStore) for tool execution * @see MCP Protocol Specification: https://modelcontextprotocol.io/specification/2025-06-18 */ export { MCPService } from './mcp.service'; @@ -286,7 +286,7 @@ export { MCPService } from './mcp.service'; * - **agenticStore**: Dispatches ToolSource.BROWSER calls here * * @see buildSandboxToolDefinition in utils/sandbox-tool - tool schema sent to the LLM - * @see agenticStore in stores/agentic.svelte.ts - tool dispatch + * @see agenticStore in stores/agentic/index.svelte.ts - tool dispatch */ export { SandboxService } from './sandbox.service'; diff --git a/tools/ui/src/lib/services/mcp.service.ts b/tools/ui/src/lib/services/mcp.service.ts index 65e9e59d6e63..7b857fd438f0 100644 --- a/tools/ui/src/lib/services/mcp.service.ts +++ b/tools/ui/src/lib/services/mcp.service.ts @@ -1,3 +1,11 @@ +/** + * MCPService - Stateless MCP protocol layer + * + * Implements the client side of the MCP spec over WebSocket, StreamableHTTP + * and SSE transports: connect, tool/prompt/resource operations and result + * formatting. No reactive state; consumed by mcpStore and its managers. + */ + import { Client } from '@modelcontextprotocol/sdk/client'; import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import { @@ -88,563 +96,149 @@ interface DiagnosticRequestDetails { export class MCPService { /** - * Create a connection log entry for phase tracking. + * Execute a tool call on a connection. + * Supports abort signal for cancellable operations (e.g., when user stops generation). + * Formats the raw tool result into a string representation. * - * @param phase - The connection phase this log belongs to - * @param message - Human-readable log message - * @param level - Log severity level (default: INFO) - * @param details - Optional structured details for debugging - * @returns Formatted connection log entry + * @param connection - The MCP connection to execute against + * @param params - Tool name and arguments to execute + * @param signal - Optional AbortSignal for cancellation support + * @returns Formatted tool execution result with content string and error flag + * @throws {Error} If tool execution fails or is aborted */ - private static createLog( - phase: MCPConnectionPhase, - message: string, - level: MCPLogLevel = MCPLogLevel.INFO, - details?: unknown - ): MCPConnectionLog { - return { - details, - level, - message, - phase, - timestamp: new Date() - }; - } + static async callTool( + connection: MCPConnection, + params: ToolCallParams, + signal?: AbortSignal + ): Promise { + throwIfAborted(signal); - private static createDiagnosticRequestDetails( - input: RequestInfo | URL, - init: RequestInit | undefined, - baseInit: RequestInit, - requestHeaders: Headers, - extraRedactedHeaders?: Iterable - ): DiagnosticRequestDetails { - const body = getRequestBody(input, init); - const details: DiagnosticRequestDetails = { - body: summarizeRequestBody(body), - credentials: init?.credentials ?? baseInit.credentials, - headers: sanitizeHeaders(requestHeaders, extraRedactedHeaders, HEADERS.PARTIAL_REDACT), - method: getRequestMethod(input, init, baseInit).toUpperCase(), - mode: init?.mode ?? baseInit.mode, - url: getRequestUrl(input) - }; - const jsonRpcMethods = extractJsonRpcMethods(body); + try { + const result = await connection.client.callTool( + { arguments: params.arguments, name: params.name }, + undefined, + { signal, timeout: connection.requestTimeoutMs } + ); - if (jsonRpcMethods) { - details.jsonRpcMethods = jsonRpcMethods; - } + return { + content: this.formatToolResult(result as ToolCallResult), + isError: (result as ToolCallResult).isError ?? false + }; + } catch (error) { + if (isAbortError(error)) { + throw error; + } - return details; - } + // Let session-expired errors propagate unwrapped for reconnection handling + if (this.isSessionExpiredError(error)) { + throw error; + } - private static addRequestHeaders( - requestHeaders: Headers, - headers: HeadersInit, - useProxy: boolean - ) { - for (const [key, value] of new Headers(headers).entries()) { - const proxiedKey = - useProxy && !key.toLowerCase().startsWith(CORS_PROXY.HEADER_PREFIX) - ? `${CORS_PROXY.HEADER_PREFIX}${key}` - : key; + const message = error instanceof Error ? error.message : String(error); - requestHeaders.set(proxiedKey, value); + throw new Error( + `Tool "${params.name}" execution failed on server "${connection.serverName}": ${message}`, + { cause: error instanceof Error ? error : undefined } + ); } } - private static summarizeError(error: unknown): Record { - if (error instanceof Error) { - return { - cause: - error.cause instanceof Error - ? { message: error.cause.message, name: error.cause.name } - : error.cause, - message: error.message, - name: error.name, - stack: error.stack?.split('\n').slice(0, 6).join('\n') - }; - } + /** + * Request completion suggestions from a server. + * Used for autocompleting prompt arguments or resource URI templates. + * + * @param connection - The MCP connection to use + * @param ref - Reference to the prompt or resource template + * @param argument - The argument being completed (name and current value) + * @returns Completion result with suggested values + */ + static async complete( + connection: MCPConnection, + ref: { type: MCPRefType.PROMPT; name: string } | { type: MCPRefType.RESOURCE; uri: string }, + argument: { name: string; value: string } + ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { + try { + const result = await connection.client.complete({ + argument, + ref + }); - return { value: String(error) }; - } + return result.completion; + } catch (error) { + console.error(`[MCPService] Failed to get completions:`, error); - private static getBrowserContext( - targetUrl: URL, - useProxy: boolean - ): Record | undefined { - if (typeof window === 'undefined') { - return undefined; + return null; } - - return { - isSecureContext: window.isSecureContext, - location: window.location.href, - origin: window.location.origin, - protocol: window.location.protocol, - sameOrigin: window.location.origin === targetUrl.origin, - targetOrigin: targetUrl.origin, - targetProtocol: targetUrl.protocol, - useProxy - }; } - private static getConnectionHints( - targetUrl: URL, - config: MCPServerConfig, - error: unknown - ): string[] { - const hints: string[] = []; - const message = error instanceof Error ? error.message : String(error); - const headerNames = Object.keys(config.headers ?? {}); - - if (typeof window !== 'undefined') { - if ( - window.location.protocol === 'https:' && - targetUrl.protocol === 'http:' && - !config.useProxy - ) { - hints.push( - 'The page is running over HTTPS but the MCP server is HTTP. Browsers often block this as mixed content; enable the proxy or use HTTPS/WSS for the MCP server.' - ); - } + /** + * Connect to a single MCP server with detailed phase tracking. + * + * Performs the full MCP connection lifecycle: + * 1. Transport creation (with automatic fallback) + * 2. Client initialization and capability exchange + * 3. Tool discovery via `listTools` + * + * Reports progress via `onPhase` callback at each step, enabling + * UI progress indicators during connection. + * + * @param serverName - Display name for the server (used in logging) + * @param serverConfig - Server URL, transport type, proxy, and auth configuration + * @param clientInfo - Optional client identification (defaults to app info) + * @param capabilities - Optional client capability declaration + * @param onPhase - Optional callback for connection phase progress updates + * @param listChangedHandlers - Optional handlers for server-initiated list change notifications + * @returns Full connection object with client, transport, tools, server info, and timing + * @throws {Error} If transport creation or connection fails + */ + static async connect( + serverName: string, + serverConfig: MCPServerConfig, + clientInfo?: Implementation, + capabilities?: ClientCapabilities, + onPhase?: MCPPhaseCallback, + listChangedHandlers?: ListChangedHandlers + ): Promise { + const startTime = performance.now(); + const effectiveClientInfo = clientInfo ?? DEFAULT_MCP_CONFIG.clientInfo; + const effectiveCapabilities = capabilities ?? DEFAULT_MCP_CONFIG.capabilities; - if (window.location.origin !== targetUrl.origin && !config.useProxy) { - hints.push( - 'This is a cross-origin browser request. If the server is reachable from curl or Node but not from the browser, missing CORS headers are the most likely cause.' - ); - } - } + // Phase: Creating transport + onPhase?.( + MCPConnectionPhase.TRANSPORT_CREATING, + this.createLog( + MCPConnectionPhase.TRANSPORT_CREATING, + `Creating transport for ${serverConfig.url}` + ) + ); - if (headerNames.length > 0) { - hints.push( - `Custom request headers are configured (${headerNames.join(', ')}). That triggers a CORS preflight, so the server must allow OPTIONS and include the matching Access-Control-Allow-Headers response.` - ); + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService][${serverName}] Creating transport...`); } - if (config.credentials && config.credentials !== 'omit') { - hints.push( - 'Credentials are enabled for this connection. Cross-origin credentialed requests need Access-Control-Allow-Credentials: true and cannot use a wildcard Access-Control-Allow-Origin.' - ); - } + const { + stopPhaseLogging, + transport, + type: transportType + } = this.createTransport(serverName, serverConfig, (log) => onPhase?.(log.phase, log)); - if (message.includes('Failed to fetch')) { - hints.push( - '"Failed to fetch" is a browser-level network failure. Common causes are CORS rejection, mixed-content blocking, certificate/TLS errors, DNS failures, or nothing listening on the target port.' - ); + // Setup WebSocket reconnection handler + if (transportType === MCPTransportType.WEBSOCKET) { + transport.onclose = () => { + console.log(`[MCPService][${serverName}] WebSocket closed, notifying for reconnection`); + onPhase?.( + MCPConnectionPhase.DISCONNECTED, + this.createLog(MCPConnectionPhase.DISCONNECTED, 'WebSocket connection closed') + ); + }; } - return hints; - } - - private static createDiagnosticFetch( - serverName: string, - config: MCPServerConfig, - baseInit: RequestInit, - targetUrl: URL, - useProxy: boolean, - onLog?: (log: MCPConnectionLog) => void - ): { - fetch: typeof fetch; - disable: () => void; - } { - let enabled = true; - - const logIfEnabled = (log: MCPConnectionLog) => { - if (enabled) { - onLog?.(log); - } - }; - - return { - disable: () => { - enabled = false; - }, - fetch: async (input, init) => { - if (useProxy && typeof window !== 'undefined') { - let requestUrlStr = ''; - - if (typeof input === 'string') { - requestUrlStr = input; - } else if (input instanceof URL) { - requestUrlStr = input.href; - } - - if (requestUrlStr) { - const parsedRequestUrl = new URL(requestUrlStr, window.location.origin); - - if ( - parsedRequestUrl.origin === window.location.origin && - !parsedRequestUrl.pathname.includes(CORS_PROXY_ENDPOINT) - ) { - const originalConfigUrl = new URL(config.url); - const realTargetUrl = new URL( - parsedRequestUrl.pathname + parsedRequestUrl.search, - originalConfigUrl.origin - ); - const proxiedUrl = buildProxiedUrl(realTargetUrl.href); - - if (typeof input === 'string') { - input = proxiedUrl.href; - } else if (input instanceof URL) { - input = proxiedUrl; - } - } - } - } - - const startedAt = performance.now(); - const requestHeaders = new Headers(baseInit.headers); - - if (typeof Request !== 'undefined' && input instanceof Request) { - this.addRequestHeaders(requestHeaders, input.headers, useProxy); - } - - if (init?.headers) { - this.addRequestHeaders(requestHeaders, init.headers, useProxy); - } - - const request = this.createDiagnosticRequestDetails( - input, - init, - baseInit, - requestHeaders, - Object.keys(config.headers ?? {}) - ); - const { method, url } = request; - - logIfEnabled( - this.createLog( - MCPConnectionPhase.INITIALIZING, - `HTTP ${method} ${url}`, - MCPLogLevel.INFO, - { - request, - serverName - } - ) - ); - - if (method === 'DELETE' && url.includes(CORS_PROXY_ENDPOINT)) { - const response = new Response(null, { status: 200, statusText: 'OK' }); - - logIfEnabled( - this.createLog( - MCPConnectionPhase.INITIALIZING, - `HTTP 200 ${method} ${url} (fake response)`, - MCPLogLevel.INFO, - { - response: { - durationMs: 0, - isFake: true, - status: response.status, - statusText: response.statusText, - url - } - } - ) - ); - - // fake response, bypass real fetch() - return response; - } - - try { - const response = await fetch(input, { - ...baseInit, - ...init, - headers: requestHeaders - }); - const durationMs = Math.round(performance.now() - startedAt); - - logIfEnabled( - this.createLog( - MCPConnectionPhase.INITIALIZING, - `HTTP ${response.status} ${method} ${url} (${durationMs}ms)`, - response.ok ? MCPLogLevel.INFO : MCPLogLevel.WARN, - { - response: { - durationMs, - headers: sanitizeHeaders(response.headers, undefined, HEADERS.PARTIAL_REDACT), - status: response.status, - statusText: response.statusText, - url - } - } - ) - ); - - return response; - } catch (error) { - const durationMs = Math.round(performance.now() - startedAt); - - logIfEnabled( - this.createLog( - MCPConnectionPhase.ERROR, - `HTTP ${method} ${url} failed: ${formatDiagnosticErrorMessage(error)}`, - MCPLogLevel.ERROR, - { - browser: this.getBrowserContext(targetUrl, useProxy), - durationMs, - error: this.summarizeError(error), - hints: this.getConnectionHints(targetUrl, config, error), - request, - serverName - } - ) - ); - - throw error; - } - } - }; - } - - /** - * Detect if an error indicates an expired/invalidated MCP session. - * Per MCP spec 2025-11-25: HTTP 404 means session invalidated, client MUST - * discard its session ID and start a new session with a fresh initialize request. - * - * @param error - The caught error to inspect - * @returns true if the error is a StreamableHTTP 404 (session not found) - */ - static isSessionExpiredError(error: unknown): boolean { - return error instanceof StreamableHTTPError && error.code === 404; - } - - /** - * Create transport based on server configuration. - * Supports WebSocket, StreamableHTTP (modern), and SSE (legacy) transports. - * When `useProxy` is enabled, routes HTTP requests through llama-server's CORS proxy. - * - * **Fallback Order:** - * 1. WebSocket — if explicitly configured (no CORS proxy support) - * 2. StreamableHTTP — default for HTTP connections - * 3. SSE — automatic fallback if StreamableHTTP fails - * - * @param config - Server configuration with url, transport type, proxy, and auth settings - * @returns Object containing the created transport and the transport type used - * @throws {Error} If url is missing, WebSocket + proxy combination, or all transports fail - */ - static createTransport( - serverName: string, - config: MCPServerConfig, - onLog?: (log: MCPConnectionLog) => void - ): { - transport: Transport; - type: MCPTransportType; - stopPhaseLogging: () => void; - } { - if (!config.url) { - throw new Error('MCP server configuration is missing url'); - } - - const useProxy = config.useProxy ?? false; - const requestInit: RequestInit = {}; - - if (config.headers) { - requestInit.headers = config.useProxy ? buildProxiedHeaders(config.headers) : config.headers; - } - - if (useProxy) { - requestInit.headers = { - ...getAuthHeaders(), - ...(requestInit.headers as Record) - }; - } - - if (config.credentials) { - requestInit.credentials = config.credentials; - } - - if (config.transport === MCPTransportType.WEBSOCKET) { - if (useProxy) { - throw new Error( - 'WebSocket transport is not supported when using CORS proxy. Use HTTP transport instead.' - ); - } - - const url = new URL(config.url); - - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService] Creating WebSocket transport for ${url.href}`); - } - - return { - stopPhaseLogging: () => {}, - transport: new WebSocketClientTransport(url), - type: MCPTransportType.WEBSOCKET - }; - } - - if (config.transport === MCPTransportType.SSE) { - const url = useProxy ? buildProxiedUrl(config.url) : new URL(config.url); - const { disable: stopPhaseLogging, fetch: diagnosticFetch } = this.createDiagnosticFetch( - serverName, - config, - requestInit, - url, - useProxy, - onLog - ); - - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService] Creating SSE transport for ${url.href}`); - } - - return { - stopPhaseLogging, - transport: new SSEClientTransport(url, { - eventSourceInit: { fetch: diagnosticFetch }, - fetch: diagnosticFetch, - requestInit - }), - type: MCPTransportType.SSE - }; - } - - const url = useProxy ? buildProxiedUrl(config.url) : new URL(config.url); - const { disable: stopPhaseLogging, fetch: diagnosticFetch } = this.createDiagnosticFetch( - serverName, - config, - requestInit, - url, - useProxy, - onLog - ); - - if (useProxy && import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService] Using CORS proxy for ${config.url} -> ${url.href}`); - } - - try { - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService] Creating StreamableHTTP transport for ${url.href}`); - } - - return { - stopPhaseLogging, - transport: new StreamableHTTPClientTransport(url, { - fetch: diagnosticFetch, - requestInit - }), - type: MCPTransportType.STREAMABLE_HTTP - }; - } catch (httpError) { - console.warn(`[MCPService] StreamableHTTP failed, trying SSE transport...`, httpError); - - try { - return { - stopPhaseLogging, - transport: new SSEClientTransport(url, { - eventSourceInit: { fetch: diagnosticFetch }, - fetch: diagnosticFetch, - requestInit - }), - type: MCPTransportType.SSE - }; - } catch (sseError) { - const httpMsg = httpError instanceof Error ? httpError.message : String(httpError); - const sseMsg = sseError instanceof Error ? sseError.message : String(sseError); - - throw new Error(`Failed to create transport. StreamableHTTP: ${httpMsg}; SSE: ${sseMsg}`); - } - } - } - - /** - * Extract server info from SDK Implementation type. - * Normalizes the SDK's server version response into our MCPServerInfo type. - * - * @param impl - Raw Implementation object from MCP SDK - * @returns Normalized server info or undefined if input is empty - */ - private static extractServerInfo(impl: Implementation | undefined): MCPServerInfo | undefined { - if (!impl) { - return undefined; - } - - return { - description: impl.description, - icons: impl.icons?.map((icon: MCPResourceIcon) => ({ - mimeType: icon.mimeType, - sizes: icon.sizes, - src: icon.src, - theme: icon.theme - })), - name: impl.name, - title: impl.title, - version: impl.version, - websiteUrl: impl.websiteUrl - }; - } - - /** - * Connect to a single MCP server with detailed phase tracking. - * - * Performs the full MCP connection lifecycle: - * 1. Transport creation (with automatic fallback) - * 2. Client initialization and capability exchange - * 3. Tool discovery via `listTools` - * - * Reports progress via `onPhase` callback at each step, enabling - * UI progress indicators during connection. - * - * @param serverName - Display name for the server (used in logging) - * @param serverConfig - Server URL, transport type, proxy, and auth configuration - * @param clientInfo - Optional client identification (defaults to app info) - * @param capabilities - Optional client capability declaration - * @param onPhase - Optional callback for connection phase progress updates - * @param listChangedHandlers - Optional handlers for server-initiated list change notifications - * @returns Full connection object with client, transport, tools, server info, and timing - * @throws {Error} If transport creation or connection fails - */ - static async connect( - serverName: string, - serverConfig: MCPServerConfig, - clientInfo?: Implementation, - capabilities?: ClientCapabilities, - onPhase?: MCPPhaseCallback, - listChangedHandlers?: ListChangedHandlers - ): Promise { - const startTime = performance.now(); - const effectiveClientInfo = clientInfo ?? DEFAULT_MCP_CONFIG.clientInfo; - const effectiveCapabilities = capabilities ?? DEFAULT_MCP_CONFIG.capabilities; - - // Phase: Creating transport - onPhase?.( - MCPConnectionPhase.TRANSPORT_CREATING, - this.createLog( - MCPConnectionPhase.TRANSPORT_CREATING, - `Creating transport for ${serverConfig.url}` - ) - ); - - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService][${serverName}] Creating transport...`); - } - - const { - stopPhaseLogging, - transport, - type: transportType - } = this.createTransport(serverName, serverConfig, (log) => onPhase?.(log.phase, log)); - - // Setup WebSocket reconnection handler - if (transportType === MCPTransportType.WEBSOCKET) { - transport.onclose = () => { - console.log(`[MCPService][${serverName}] WebSocket closed, notifying for reconnection`); - onPhase?.( - MCPConnectionPhase.DISCONNECTED, - this.createLog(MCPConnectionPhase.DISCONNECTED, 'WebSocket connection closed') - ); - }; - } - - // Phase: Transport ready - onPhase?.( - MCPConnectionPhase.TRANSPORT_READY, - this.createLog(MCPConnectionPhase.TRANSPORT_READY, `Transport ready (${transportType})`), - { transportType } - ); + // Phase: Transport ready + onPhase?.( + MCPConnectionPhase.TRANSPORT_READY, + this.createLog(MCPConnectionPhase.TRANSPORT_READY, `Transport ready (${transportType})`), + { transportType } + ); const client = new Client( { @@ -848,36 +442,322 @@ export class MCPService { } /** - * Disconnect from a server. - * Clears the `onclose` handler to prevent reconnection attempts on voluntary disconnect. + * Create transport based on server configuration. + * Supports WebSocket, StreamableHTTP (modern), and SSE (legacy) transports. + * When `useProxy` is enabled, routes HTTP requests through llama-server's CORS proxy. + * + * **Fallback Order:** + * 1. WebSocket — if explicitly configured (no CORS proxy support) + * 2. StreamableHTTP — default for HTTP connections + * 3. SSE — automatic fallback if StreamableHTTP fails + * + * @param config - Server configuration with url, transport type, proxy, and auth settings + * @returns Object containing the created transport and the transport type used + * @throws {Error} If url is missing, WebSocket + proxy combination, or all transports fail + */ + static createTransport( + serverName: string, + config: MCPServerConfig, + onLog?: (log: MCPConnectionLog) => void + ): { + transport: Transport; + type: MCPTransportType; + stopPhaseLogging: () => void; + } { + if (!config.url) { + throw new Error('MCP server configuration is missing url'); + } + + const useProxy = config.useProxy ?? false; + const requestInit: RequestInit = {}; + + if (config.headers) { + requestInit.headers = config.useProxy ? buildProxiedHeaders(config.headers) : config.headers; + } + + if (useProxy) { + requestInit.headers = { + ...getAuthHeaders(), + ...(requestInit.headers as Record) + }; + } + + if (config.credentials) { + requestInit.credentials = config.credentials; + } + + if (config.transport === MCPTransportType.WEBSOCKET) { + if (useProxy) { + throw new Error( + 'WebSocket transport is not supported when using CORS proxy. Use HTTP transport instead.' + ); + } + + const url = new URL(config.url); + + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService] Creating WebSocket transport for ${url.href}`); + } + + return { + stopPhaseLogging: () => {}, + transport: new WebSocketClientTransport(url), + type: MCPTransportType.WEBSOCKET + }; + } + + if (config.transport === MCPTransportType.SSE) { + const url = useProxy ? buildProxiedUrl(config.url) : new URL(config.url); + const { disable: stopPhaseLogging, fetch: diagnosticFetch } = this.createDiagnosticFetch( + serverName, + config, + requestInit, + url, + useProxy, + onLog + ); + + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService] Creating SSE transport for ${url.href}`); + } + + return { + stopPhaseLogging, + transport: new SSEClientTransport(url, { + eventSourceInit: { fetch: diagnosticFetch }, + fetch: diagnosticFetch, + requestInit + }), + type: MCPTransportType.SSE + }; + } + + const url = useProxy ? buildProxiedUrl(config.url) : new URL(config.url); + const { disable: stopPhaseLogging, fetch: diagnosticFetch } = this.createDiagnosticFetch( + serverName, + config, + requestInit, + url, + useProxy, + onLog + ); + + if (useProxy && import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService] Using CORS proxy for ${config.url} -> ${url.href}`); + } + + try { + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService] Creating StreamableHTTP transport for ${url.href}`); + } + + return { + stopPhaseLogging, + transport: new StreamableHTTPClientTransport(url, { + fetch: diagnosticFetch, + requestInit + }), + type: MCPTransportType.STREAMABLE_HTTP + }; + } catch (httpError) { + console.warn(`[MCPService] StreamableHTTP failed, trying SSE transport...`, httpError); + + try { + return { + stopPhaseLogging, + transport: new SSEClientTransport(url, { + eventSourceInit: { fetch: diagnosticFetch }, + fetch: diagnosticFetch, + requestInit + }), + type: MCPTransportType.SSE + }; + } catch (sseError) { + const httpMsg = httpError instanceof Error ? httpError.message : String(httpError); + const sseMsg = sseError instanceof Error ? sseError.message : String(sseError); + + throw new Error(`Failed to create transport. StreamableHTTP: ${httpMsg}; SSE: ${sseMsg}`); + } + } + } + + /** + * Disconnect from a server. + * Clears the `onclose` handler to prevent reconnection attempts on voluntary disconnect. + * + * @param connection - The active MCP connection to close + */ + static async disconnect(connection: MCPConnection): Promise { + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService][${connection.serverName}] Disconnecting...`); + } + + try { + // Terminate the session first for streamable-http transports to cleanly + // close streams, matching the inspector's disconnect flow. + if (connection.transport instanceof StreamableHTTPClientTransport) { + await connection.transport.terminateSession(); + } + + // Clear error handlers before closing to prevent noise from expected + // abort errors during shutdown. The inspector avoids this entirely + // by not setting onerror, but since we use it for protocol logging, + // we must clear it before disconnect. + connection.client.onerror = undefined; + + if (connection.transport.onclose) { + connection.transport.onclose = undefined; + } + + await connection.client.close(); + } catch (error) { + console.warn(`[MCPService][${connection.serverName}] Error during disconnect:`, error); + } + } + + /** + * Get a specific prompt with arguments. + * Unlike list operations, this throws on failure since the caller explicitly + * requested a specific prompt and needs to handle the error. + * + * @param connection - The MCP connection to use + * @param name - The prompt name to retrieve + * @param args - Optional key-value arguments to pass to the prompt + * @returns The prompt result with messages and metadata + * @throws {Error} If the prompt retrieval fails + */ + static async getPrompt( + connection: MCPConnection, + name: string, + args?: Record + ): Promise { + try { + return await connection.client.getPrompt({ arguments: args, name }); + } catch (error) { + console.error(`[MCPService][${connection.serverName}] Failed to get prompt:`, error); + + throw error; + } + } + + /** + * Detect if an error indicates an expired/invalidated MCP session. + * Per MCP spec 2025-11-25: HTTP 404 means session invalidated, client MUST + * discard its session ID and start a new session with a fresh initialize request. + * + * @param error - The caught error to inspect + * @returns true if the error is a StreamableHTTP 404 (session not found) + */ + static isSessionExpiredError(error: unknown): boolean { + return error instanceof StreamableHTTPError && error.code === 404; + } + + /** + * List all resources from a connection (handles pagination automatically). + * @param connection - The MCP connection to use + * @returns Array of all available resources + */ + static async listAllResources(connection: MCPConnection): Promise { + return this.paginate( + connection, + (cursor) => this.listResources(connection, cursor), + (result) => result.resources + ); + } + + /** + * List all resource templates from a connection (handles pagination automatically). + * @param connection - The MCP connection to use + * @returns Array of all available resource templates + */ + static async listAllResourceTemplates(connection: MCPConnection): Promise { + return this.paginate( + connection, + (cursor) => this.listResourceTemplates(connection, cursor), + (result) => result.resourceTemplates + ); + } + + /** + * List prompts from a connection. + * Silently returns empty array on failure (logged as warning). * - * @param connection - The active MCP connection to close + * @param connection - The MCP connection to query + * @returns Array of available prompts, or empty array on error */ - static async disconnect(connection: MCPConnection): Promise { - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService][${connection.serverName}] Disconnecting...`); + static async listPrompts(connection: MCPConnection): Promise { + try { + const result = await connection.client.listPrompts(); + + return result.prompts ?? []; + } catch (error) { + // Let session-expired errors propagate for reconnection handling + if (this.isSessionExpiredError(error)) { + throw error; + } + + console.warn(`[MCPService][${connection.serverName}] Failed to list prompts:`, error); + + return []; } + } + /** + * List resources from a connection. + * @param connection - The MCP connection to use + * @param cursor - Optional pagination cursor + * @returns Array of available resources and optional next cursor + */ + static async listResources( + connection: MCPConnection, + cursor?: string + ): Promise<{ resources: MCPResource[]; nextCursor?: string }> { try { - // Terminate the session first for streamable-http transports to cleanly - // close streams, matching the inspector's disconnect flow. - if (connection.transport instanceof StreamableHTTPClientTransport) { - await connection.transport.terminateSession(); + const result = await connection.client.listResources(cursor ? { cursor } : undefined); + + return { + nextCursor: result.nextCursor, + resources: (result.resources ?? []) as MCPResource[] + }; + } catch (error) { + if (this.isSessionExpiredError(error)) { + throw error; } - // Clear error handlers before closing to prevent noise from expected - // abort errors during shutdown. The inspector avoids this entirely - // by not setting onerror, but since we use it for protocol logging, - // we must clear it before disconnect. - connection.client.onerror = undefined; + console.warn(`[MCPService][${connection.serverName}] Failed to list resources:`, error); - if (connection.transport.onclose) { - connection.transport.onclose = undefined; - } + return { resources: [] }; + } + } - await connection.client.close(); + /** + * List resource templates from a connection. + * @param connection - The MCP connection to use + * @param cursor - Optional pagination cursor + * @returns Array of available resource templates and optional next cursor + */ + static async listResourceTemplates( + connection: MCPConnection, + cursor?: string + ): Promise<{ resourceTemplates: MCPResourceTemplate[]; nextCursor?: string }> { + try { + const result = await connection.client.listResourceTemplates(cursor ? { cursor } : undefined); + + return { + nextCursor: result.nextCursor, + resourceTemplates: (result.resourceTemplates ?? []) as MCPResourceTemplate[] + }; } catch (error) { - console.warn(`[MCPService][${connection.serverName}] Error during disconnect:`, error); + if (this.isSessionExpiredError(error)) { + throw error; + } + + console.warn( + `[MCPService][${connection.serverName}] Failed to list resource templates:`, + error + ); + + return { resourceTemplates: [] }; } } @@ -906,386 +786,506 @@ export class MCPService { } /** - * List prompts from a connection. - * Silently returns empty array on failure (logged as warning). - * - * @param connection - The MCP connection to query - * @returns Array of available prompts, or empty array on error + * Read the contents of a resource. + * @param connection - The MCP connection to use + * @param uri - The URI of the resource to read + * @returns The resource contents */ - static async listPrompts(connection: MCPConnection): Promise { + static async readResource( + connection: MCPConnection, + uri: string + ): Promise { try { - const result = await connection.client.listPrompts(); + const result = await connection.client.readResource({ uri }); - return result.prompts ?? []; + return { + _meta: result._meta, + contents: (result.contents ?? []) as MCPResourceContent[] + }; } catch (error) { - // Let session-expired errors propagate for reconnection handling - if (this.isSessionExpiredError(error)) { - throw error; - } - - console.warn(`[MCPService][${connection.serverName}] Failed to list prompts:`, error); + console.error(`[MCPService][${connection.serverName}] Failed to read resource:`, error); - return []; + throw error; } } /** - * Get a specific prompt with arguments. - * Unlike list operations, this throws on failure since the caller explicitly - * requested a specific prompt and needs to handle the error. - * + * Subscribe to updates for a resource. + * The server will send notifications/resources/updated when the resource changes. * @param connection - The MCP connection to use - * @param name - The prompt name to retrieve - * @param args - Optional key-value arguments to pass to the prompt - * @returns The prompt result with messages and metadata - * @throws {Error} If the prompt retrieval fails + * @param uri - The URI of the resource to subscribe to */ - static async getPrompt( - connection: MCPConnection, - name: string, - args?: Record - ): Promise { + static async subscribeResource(connection: MCPConnection, uri: string): Promise { try { - return await connection.client.getPrompt({ arguments: args, name }); + await connection.client.subscribeResource({ uri }); + + console.log(`[MCPService][${connection.serverName}] Subscribed to resource: ${uri}`); } catch (error) { - console.error(`[MCPService][${connection.serverName}] Failed to get prompt:`, error); + console.error( + `[MCPService][${connection.serverName}] Failed to subscribe to resource:`, + error + ); throw error; } } /** - * Execute a tool call on a connection. - * Supports abort signal for cancellable operations (e.g., when user stops generation). - * Formats the raw tool result into a string representation. + * Check if a connection supports resources. + * Per MCP spec: presence of the `resources` key (even as empty object `{}`) indicates support. + * Empty object means resources are supported but no sub-features (subscribe, listChanged). * - * @param connection - The MCP connection to execute against - * @param params - Tool name and arguments to execute - * @param signal - Optional AbortSignal for cancellation support - * @returns Formatted tool execution result with content string and error flag - * @throws {Error} If tool execution fails or is aborted + * @param connection - The MCP connection to check + * @returns Whether the server declares the resources capability */ - static async callTool( - connection: MCPConnection, - params: ToolCallParams, - signal?: AbortSignal - ): Promise { - throwIfAborted(signal); + static supportsResources(connection: MCPConnection): boolean { + // Per MCP spec: "Servers that support resources MUST declare the resources capability" + // The presence of the key indicates support, even if it's an empty object + return connection.serverCapabilities?.resources !== undefined; + } + + /** + * Check if a connection supports resource subscriptions. + * @param connection - The MCP connection to check + * @returns Whether the server supports resource subscriptions + */ + static supportsResourceSubscriptions(connection: MCPConnection): boolean { + return !!connection.serverCapabilities?.resources?.subscribe; + } + /** + * Unsubscribe from updates for a resource. + * @param connection - The MCP connection to use + * @param uri - The URI of the resource to unsubscribe from + */ + static async unsubscribeResource(connection: MCPConnection, uri: string): Promise { try { - const result = await connection.client.callTool( - { arguments: params.arguments, name: params.name }, - undefined, - { signal, timeout: connection.requestTimeoutMs } + await connection.client.unsubscribeResource({ uri }); + + if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { + console.log(`[MCPService][${connection.serverName}] Unsubscribed from resource: ${uri}`); + } + } catch (error) { + console.error( + `[MCPService][${connection.serverName}] Failed to unsubscribe from resource:`, + error ); - return { - content: this.formatToolResult(result as ToolCallResult), - isError: (result as ToolCallResult).isError ?? false - }; - } catch (error) { - if (isAbortError(error)) { - throw error; - } + throw error; + } + } + + private static addRequestHeaders( + requestHeaders: Headers, + headers: HeadersInit, + useProxy: boolean + ) { + for (const [key, value] of new Headers(headers).entries()) { + const proxiedKey = + useProxy && !key.toLowerCase().startsWith(CORS_PROXY.HEADER_PREFIX) + ? `${CORS_PROXY.HEADER_PREFIX}${key}` + : key; + + requestHeaders.set(proxiedKey, value); + } + } + + private static createDiagnosticFetch( + serverName: string, + config: MCPServerConfig, + baseInit: RequestInit, + targetUrl: URL, + useProxy: boolean, + onLog?: (log: MCPConnectionLog) => void + ): { + fetch: typeof fetch; + disable: () => void; + } { + let enabled = true; + + const logIfEnabled = (log: MCPConnectionLog) => { + if (enabled) { + onLog?.(log); + } + }; + + return { + disable: () => { + enabled = false; + }, + fetch: async (input, init) => { + if (useProxy && typeof window !== 'undefined') { + let requestUrlStr = ''; + + if (typeof input === 'string') { + requestUrlStr = input; + } else if (input instanceof URL) { + requestUrlStr = input.href; + } + + if (requestUrlStr) { + const parsedRequestUrl = new URL(requestUrlStr, window.location.origin); + + if ( + parsedRequestUrl.origin === window.location.origin && + !parsedRequestUrl.pathname.includes(CORS_PROXY_ENDPOINT) + ) { + const originalConfigUrl = new URL(config.url); + const realTargetUrl = new URL( + parsedRequestUrl.pathname + parsedRequestUrl.search, + originalConfigUrl.origin + ); + const proxiedUrl = buildProxiedUrl(realTargetUrl.href); - // Let session-expired errors propagate unwrapped for reconnection handling - if (this.isSessionExpiredError(error)) { - throw error; - } + if (typeof input === 'string') { + input = proxiedUrl.href; + } else if (input instanceof URL) { + input = proxiedUrl; + } + } + } + } - const message = error instanceof Error ? error.message : String(error); + const startedAt = performance.now(); + const requestHeaders = new Headers(baseInit.headers); - throw new Error( - `Tool "${params.name}" execution failed on server "${connection.serverName}": ${message}`, - { cause: error instanceof Error ? error : undefined } - ); - } - } + if (typeof Request !== 'undefined' && input instanceof Request) { + this.addRequestHeaders(requestHeaders, input.headers, useProxy); + } - /** - * Format tool result content items to a single string. - * Handles text, image (base64 data URL), and embedded resource content types. - * - * @param result - Raw tool call result from MCP SDK - * @returns Concatenated string representation of all content items - */ - private static formatToolResult(result: ToolCallResult): string { - const content = result.content; + if (init?.headers) { + this.addRequestHeaders(requestHeaders, init.headers, useProxy); + } - if (!Array.isArray(content)) return ''; + const request = this.createDiagnosticRequestDetails( + input, + init, + baseInit, + requestHeaders, + Object.keys(config.headers ?? {}) + ); + const { method, url } = request; - const formatted = content - .map((item) => this.formatSingleContent(item)) - .filter(Boolean) - .join(NEWLINE); + logIfEnabled( + this.createLog( + MCPConnectionPhase.INITIALIZING, + `HTTP ${method} ${url}`, + MCPLogLevel.INFO, + { + request, + serverName + } + ) + ); - if (formatted !== '') { - return formatted; - } + if (method === 'DELETE' && url.includes(CORS_PROXY_ENDPOINT)) { + const response = new Response(null, { status: 200, statusText: 'OK' }); - if (result.structuredContent && typeof result.structuredContent === 'object') { - return JSON.stringify(result.structuredContent); - } + logIfEnabled( + this.createLog( + MCPConnectionPhase.INITIALIZING, + `HTTP 200 ${method} ${url} (fake response)`, + MCPLogLevel.INFO, + { + response: { + durationMs: 0, + isFake: true, + status: response.status, + statusText: response.statusText, + url + } + } + ) + ); - return ''; - } + // fake response, bypass real fetch() + return response; + } - private static formatSingleContent(content: ToolResultContentItem): string { - if (content.type === MCPContentType.TEXT && content.text) { - return content.text; - } + try { + const response = await fetch(input, { + ...baseInit, + ...init, + headers: requestHeaders + }); + const durationMs = Math.round(performance.now() - startedAt); - if (content.type === MCPContentType.IMAGE && content.data) { - return createBase64DataUrl(content.mimeType ?? DEFAULT_IMAGE_MIME_TYPE, content.data); - } + logIfEnabled( + this.createLog( + MCPConnectionPhase.INITIALIZING, + `HTTP ${response.status} ${method} ${url} (${durationMs}ms)`, + response.ok ? MCPLogLevel.INFO : MCPLogLevel.WARN, + { + response: { + durationMs, + headers: sanitizeHeaders(response.headers, undefined, HEADERS.PARTIAL_REDACT), + status: response.status, + statusText: response.statusText, + url + } + } + ) + ); - if (content.type === MCPContentType.RESOURCE && content.resource) { - const resource = content.resource; + return response; + } catch (error) { + const durationMs = Math.round(performance.now() - startedAt); - if (resource.text) return resource.text; + logIfEnabled( + this.createLog( + MCPConnectionPhase.ERROR, + `HTTP ${method} ${url} failed: ${formatDiagnosticErrorMessage(error)}`, + MCPLogLevel.ERROR, + { + browser: this.getBrowserContext(targetUrl, useProxy), + durationMs, + error: this.summarizeError(error), + hints: this.getConnectionHints(targetUrl, config, error), + request, + serverName + } + ) + ); - if (resource.blob) return resource.blob; + throw error; + } + } + }; + } - return JSON.stringify(resource); - } + private static createDiagnosticRequestDetails( + input: RequestInfo | URL, + init: RequestInit | undefined, + baseInit: RequestInit, + requestHeaders: Headers, + extraRedactedHeaders?: Iterable + ): DiagnosticRequestDetails { + const body = getRequestBody(input, init); + const details: DiagnosticRequestDetails = { + body: summarizeRequestBody(body), + credentials: init?.credentials ?? baseInit.credentials, + headers: sanitizeHeaders(requestHeaders, extraRedactedHeaders, HEADERS.PARTIAL_REDACT), + method: getRequestMethod(input, init, baseInit).toUpperCase(), + mode: init?.mode ?? baseInit.mode, + url: getRequestUrl(input) + }; + const jsonRpcMethods = extractJsonRpcMethods(body); - if (content.data && content.mimeType) { - return createBase64DataUrl(content.mimeType, content.data); + if (jsonRpcMethods) { + details.jsonRpcMethods = jsonRpcMethods; } - return JSON.stringify(content); + return details; } /** + * Create a connection log entry for phase tracking. * - * - * Completions Operations - * - * + * @param phase - The connection phase this log belongs to + * @param message - Human-readable log message + * @param level - Log severity level (default: INFO) + * @param details - Optional structured details for debugging + * @returns Formatted connection log entry */ + private static createLog( + phase: MCPConnectionPhase, + message: string, + level: MCPLogLevel = MCPLogLevel.INFO, + details?: unknown + ): MCPConnectionLog { + return { + details, + level, + message, + phase, + timestamp: new Date() + }; + } /** - * Request completion suggestions from a server. - * Used for autocompleting prompt arguments or resource URI templates. + * Extract server info from SDK Implementation type. + * Normalizes the SDK's server version response into our MCPServerInfo type. * - * @param connection - The MCP connection to use - * @param ref - Reference to the prompt or resource template - * @param argument - The argument being completed (name and current value) - * @returns Completion result with suggested values + * @param impl - Raw Implementation object from MCP SDK + * @returns Normalized server info or undefined if input is empty */ - static async complete( - connection: MCPConnection, - ref: { type: MCPRefType.PROMPT; name: string } | { type: MCPRefType.RESOURCE; uri: string }, - argument: { name: string; value: string } - ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { - try { - const result = await connection.client.complete({ - argument, - ref - }); + private static extractServerInfo(impl: Implementation | undefined): MCPServerInfo | undefined { + if (!impl) { + return undefined; + } - return result.completion; - } catch (error) { - console.error(`[MCPService] Failed to get completions:`, error); + return { + description: impl.description, + icons: impl.icons?.map((icon: MCPResourceIcon) => ({ + mimeType: icon.mimeType, + sizes: icon.sizes, + src: icon.src, + theme: icon.theme + })), + name: impl.name, + title: impl.title, + version: impl.version, + websiteUrl: impl.websiteUrl + }; + } + + private static formatSingleContent(content: ToolResultContentItem): string { + if (content.type === MCPContentType.TEXT && content.text) { + return content.text; + } - return null; + if (content.type === MCPContentType.IMAGE && content.data) { + return createBase64DataUrl(content.mimeType ?? DEFAULT_IMAGE_MIME_TYPE, content.data); } - } - /** - * - * - * Resources Operations - * - * - */ + if (content.type === MCPContentType.RESOURCE && content.resource) { + const resource = content.resource; - /** - * List resources from a connection. - * @param connection - The MCP connection to use - * @param cursor - Optional pagination cursor - * @returns Array of available resources and optional next cursor - */ - static async listResources( - connection: MCPConnection, - cursor?: string - ): Promise<{ resources: MCPResource[]; nextCursor?: string }> { - try { - const result = await connection.client.listResources(cursor ? { cursor } : undefined); + if (resource.text) return resource.text; - return { - nextCursor: result.nextCursor, - resources: (result.resources ?? []) as MCPResource[] - }; - } catch (error) { - if (this.isSessionExpiredError(error)) { - throw error; - } + if (resource.blob) return resource.blob; - console.warn(`[MCPService][${connection.serverName}] Failed to list resources:`, error); + return JSON.stringify(resource); + } - return { resources: [] }; + if (content.data && content.mimeType) { + return createBase64DataUrl(content.mimeType, content.data); } + + return JSON.stringify(content); } /** - * List all resources from a connection (handles pagination automatically). - * @param connection - The MCP connection to use - * @returns Array of all available resources + * Format tool result content items to a single string. + * Handles text, image (base64 data URL), and embedded resource content types. + * + * @param result - Raw tool call result from MCP SDK + * @returns Concatenated string representation of all content items */ - static async listAllResources(connection: MCPConnection): Promise { - const allResources: MCPResource[] = []; + private static formatToolResult(result: ToolCallResult): string { + const content = result.content; - let cursor: string | undefined; + if (!Array.isArray(content)) return ''; - do { - const result = await this.listResources(connection, cursor); + const formatted = content + .map((item) => this.formatSingleContent(item)) + .filter(Boolean) + .join(NEWLINE); - allResources.push(...result.resources); - cursor = result.nextCursor; - } while (cursor); + if (formatted !== '') { + return formatted; + } + + if (result.structuredContent && typeof result.structuredContent === 'object') { + return JSON.stringify(result.structuredContent); + } - return allResources; + return ''; } - /** - * List resource templates from a connection. - * @param connection - The MCP connection to use - * @param cursor - Optional pagination cursor - * @returns Array of available resource templates and optional next cursor - */ - static async listResourceTemplates( - connection: MCPConnection, - cursor?: string - ): Promise<{ resourceTemplates: MCPResourceTemplate[]; nextCursor?: string }> { - try { - const result = await connection.client.listResourceTemplates(cursor ? { cursor } : undefined); + private static getBrowserContext( + targetUrl: URL, + useProxy: boolean + ): Record | undefined { + if (typeof window === 'undefined') { + return undefined; + } - return { - nextCursor: result.nextCursor, - resourceTemplates: (result.resourceTemplates ?? []) as MCPResourceTemplate[] - }; - } catch (error) { - if (this.isSessionExpiredError(error)) { - throw error; + return { + isSecureContext: window.isSecureContext, + location: window.location.href, + origin: window.location.origin, + protocol: window.location.protocol, + sameOrigin: window.location.origin === targetUrl.origin, + targetOrigin: targetUrl.origin, + targetProtocol: targetUrl.protocol, + useProxy + }; + } + + private static getConnectionHints( + targetUrl: URL, + config: MCPServerConfig, + error: unknown + ): string[] { + const hints: string[] = []; + const message = error instanceof Error ? error.message : String(error); + const headerNames = Object.keys(config.headers ?? {}); + + if (typeof window !== 'undefined') { + if ( + window.location.protocol === 'https:' && + targetUrl.protocol === 'http:' && + !config.useProxy + ) { + hints.push( + 'The page is running over HTTPS but the MCP server is HTTP. Browsers often block this as mixed content; enable the proxy or use HTTPS/WSS for the MCP server.' + ); } - console.warn( - `[MCPService][${connection.serverName}] Failed to list resource templates:`, - error + if (window.location.origin !== targetUrl.origin && !config.useProxy) { + hints.push( + 'This is a cross-origin browser request. If the server is reachable from curl or Node but not from the browser, missing CORS headers are the most likely cause.' + ); + } + } + + if (headerNames.length > 0) { + hints.push( + `Custom request headers are configured (${headerNames.join(', ')}). That triggers a CORS preflight, so the server must allow OPTIONS and include the matching Access-Control-Allow-Headers response.` + ); + } + + if (config.credentials && config.credentials !== 'omit') { + hints.push( + 'Credentials are enabled for this connection. Cross-origin credentialed requests need Access-Control-Allow-Credentials: true and cannot use a wildcard Access-Control-Allow-Origin.' ); + } - return { resourceTemplates: [] }; + if (message.includes('Failed to fetch')) { + hints.push( + '"Failed to fetch" is a browser-level network failure. Common causes are CORS rejection, mixed-content blocking, certificate/TLS errors, DNS failures, or nothing listening on the target port.' + ); } + + return hints; } /** - * List all resource templates from a connection (handles pagination automatically). - * @param connection - The MCP connection to use - * @returns Array of all available resource templates + * Walk a cursor-paginated MCP list endpoint, collecting every page. */ - static async listAllResourceTemplates(connection: MCPConnection): Promise { - const allTemplates: MCPResourceTemplate[] = []; + private static async paginate( + connection: MCPConnection, + fetchPage: (cursor?: string) => Promise, + extract: (result: R) => T[] + ): Promise { + const all: T[] = []; let cursor: string | undefined; do { - const result = await this.listResourceTemplates(connection, cursor); + const result = await fetchPage(cursor); - allTemplates.push(...result.resourceTemplates); + all.push(...extract(result)); cursor = result.nextCursor; } while (cursor); - return allTemplates; + return all; } - /** - * Read the contents of a resource. - * @param connection - The MCP connection to use - * @param uri - The URI of the resource to read - * @returns The resource contents - */ - static async readResource( - connection: MCPConnection, - uri: string - ): Promise { - try { - const result = await connection.client.readResource({ uri }); - + private static summarizeError(error: unknown): Record { + if (error instanceof Error) { return { - _meta: result._meta, - contents: (result.contents ?? []) as MCPResourceContent[] + cause: + error.cause instanceof Error + ? { message: error.cause.message, name: error.cause.name } + : error.cause, + message: error.message, + name: error.name, + stack: error.stack?.split('\n').slice(0, 6).join('\n') }; - } catch (error) { - console.error(`[MCPService][${connection.serverName}] Failed to read resource:`, error); - - throw error; - } - } - - /** - * Subscribe to updates for a resource. - * The server will send notifications/resources/updated when the resource changes. - * @param connection - The MCP connection to use - * @param uri - The URI of the resource to subscribe to - */ - static async subscribeResource(connection: MCPConnection, uri: string): Promise { - try { - await connection.client.subscribeResource({ uri }); - - console.log(`[MCPService][${connection.serverName}] Subscribed to resource: ${uri}`); - } catch (error) { - console.error( - `[MCPService][${connection.serverName}] Failed to subscribe to resource:`, - error - ); - - throw error; - } - } - - /** - * Unsubscribe from updates for a resource. - * @param connection - The MCP connection to use - * @param uri - The URI of the resource to unsubscribe from - */ - static async unsubscribeResource(connection: MCPConnection, uri: string): Promise { - try { - await connection.client.unsubscribeResource({ uri }); - - if (import.meta.env.DEV && import.meta.env.VITE_DEBUG) { - console.log(`[MCPService][${connection.serverName}] Unsubscribed from resource: ${uri}`); - } - } catch (error) { - console.error( - `[MCPService][${connection.serverName}] Failed to unsubscribe from resource:`, - error - ); - - throw error; } - } - - /** - * Check if a connection supports resources. - * Per MCP spec: presence of the `resources` key (even as empty object `{}`) indicates support. - * Empty object means resources are supported but no sub-features (subscribe, listChanged). - * - * @param connection - The MCP connection to check - * @returns Whether the server declares the resources capability - */ - static supportsResources(connection: MCPConnection): boolean { - // Per MCP spec: "Servers that support resources MUST declare the resources capability" - // The presence of the key indicates support, even if it's an empty object - return connection.serverCapabilities?.resources !== undefined; - } - /** - * Check if a connection supports resource subscriptions. - * @param connection - The MCP connection to check - * @returns Whether the server supports resource subscriptions - */ - static supportsResourceSubscriptions(connection: MCPConnection): boolean { - return !!connection.serverCapabilities?.resources?.subscribe; + return { value: String(error) }; } } diff --git a/tools/ui/src/lib/services/migration.service.ts b/tools/ui/src/lib/services/migration.service.ts index 2626a42b3c99..5d321b3ba284 100644 --- a/tools/ui/src/lib/services/migration.service.ts +++ b/tools/ui/src/lib/services/migration.service.ts @@ -1,20 +1,11 @@ /** - * Migration Service - Unified data migration hook + * MigrationService - Unified data migration hook * - * Centralizes all data migrations (localStorage, IndexedDB, legacy formats) into a single - * initialization point. Each migration copies data to new format WITHOUT deleting the old. - * - * **Architecture:** - * - Migrations are defined as objects with `id` and `run()` methods - * - Migration state is tracked in localStorage to avoid re-running - * - `runAllMigrations()` should be called once at app startup - * - All migrations are NON-DESTRUCTIVE - legacy data is preserved for downgrade compatibility - * - * **Current Migrations:** - * 1. localStorage prefix: Copy LlamaCppWebui.* → LlamaUi.* (both preserved) - * 2. IndexedDB database: Copy LlamacppWebui → LlamaUi (both preserved) - * 3. Legacy message format: Transform in-place (preserves structure, migrates markers) - * 4. Theme key: Copy standalone `theme` → config object (both preserved) + * Centralizes all data migrations (localStorage, IndexedDB, legacy formats) + * into a single initialization point. Each migration copies data to the new + * format WITHOUT deleting the old, and state is tracked in localStorage so + * `runAllMigrations()` (called once at startup) never re-runs a completed + * migration. All migrations are non-destructive for downgrade compatibility. */ import { diff --git a/tools/ui/src/lib/services/models.service.ts b/tools/ui/src/lib/services/models.service.ts index 84832e086f5e..bb1bbd356a71 100644 --- a/tools/ui/src/lib/services/models.service.ts +++ b/tools/ui/src/lib/services/models.service.ts @@ -1,24 +1,54 @@ +/** + * ModelsService - Stateless model management API layer + * + * Wraps the /models endpoints (list, load, unload) and the /models/sse + * status feed in MODEL and ROUTER modes. No reactive state; consumed by + * modelsStore and its status manager. + */ + import { base } from '$app/paths'; -import { - API_MODELS, - MODEL_ID, - SSE_DATA_PREFIX, - SSE_LINE_SEPARATOR, - SSE_RECORD_SEPARATOR -} from '$lib/constants'; +import { API_MODELS, MODEL_ID } from '$lib/constants'; import { ServerModelStatus } from '$lib/enums'; import type { ParsedModelId } from '$lib/types/models'; -import { apiFetch, apiPost, normalizeModelName } from '$lib/utils'; +import { + apiFetch, + apiPost, + extractSseDataPayload, + normalizeModelName, + splitSseRecords +} from '$lib/utils'; import { getAuthHeaders } from '$lib/utils/api-headers'; export class ModelsService { + private static readonly SSE_RECONNECT_MS = 1000; + /** + * Check if a model is loaded based on its metadata. * + * @param model - Model data entry from the API response + * @returns True if the model status is LOADED + */ + static isModelLoaded(model: ApiModelDataEntry): boolean { + return model.status.value === ServerModelStatus.LOADED; + } + + /** * - * Listing * + * Load/Unload * + * + */ + + /** + * Check if a model is currently loading. + * + * @param model - Model data entry from the API response + * @returns True if the model status is LOADING */ + static isModelLoading(model: ApiModelDataEntry): boolean { + return model.status.value === ServerModelStatus.LOADING; + } /** * Fetch list of models from OpenAI-compatible endpoint. @@ -41,14 +71,6 @@ export class ModelsService { return apiFetch(API_MODELS.LIST); } - /** - * - * - * Load/Unload - * - * - */ - /** * Load a model (ROUTER mode only). * Sends POST request to `/models/load`. Note: the endpoint returns success @@ -68,137 +90,6 @@ export class ModelsService { return apiPost(API_MODELS.LOAD, payload); } - /** - * Unload a model (ROUTER mode only). - * Sends POST request to `/models/unload`. Note: the endpoint returns success - * before unloading completes — use polling to await actual unload status. - * - * @param modelId - Model identifier to unload - * @returns Unload response from the server - */ - static async unload(modelId: string): Promise { - return apiPost(API_MODELS.UNLOAD, { model: modelId }); - } - - /** - * - * - * Status - * - * - */ - - /** - * Check if a model is loaded based on its metadata. - * - * @param model - Model data entry from the API response - * @returns True if the model status is LOADED - */ - static isModelLoaded(model: ApiModelDataEntry): boolean { - return model.status.value === ServerModelStatus.LOADED; - } - - /** - * Check if a model is currently loading. - * - * @param model - Model data entry from the API response - * @returns True if the model status is LOADING - */ - static isModelLoading(model: ApiModelDataEntry): boolean { - return model.status.value === ServerModelStatus.LOADING; - } - - /** - * - * - * Status Feed - * - * - */ - - private static readonly SSE_RECONNECT_MS = 1000; - - /** - * Read the /models/sse feed and invoke onEvent for each parsed envelope. - * Reconnects on network drops until the signal aborts. Splits the byte - * stream into SSE records on the blank line boundary; the payload rides in - * the data lines as a JSON envelope with its own model, event and data fields. - */ - static async watchModelEvents( - signal: AbortSignal, - onEvent: (event: ApiModelsSseEvent) => void - ): Promise { - const decoder = new TextDecoder(); - - while (!signal.aborted) { - try { - const response = await fetch(`${base}${API_MODELS.SSE}`, { - headers: getAuthHeaders(), - signal - }); - - if (response.ok && response.body) { - const reader = response.body.getReader(); - - let buffer = ''; - - while (!signal.aborted) { - const { done, value } = await reader.read(); - - if (done) break; - - buffer += decoder.decode(value, { stream: true }); - - let boundary = buffer.indexOf(SSE_RECORD_SEPARATOR); - - while (boundary !== -1) { - const event = ModelsService.parseStatusRecord(buffer.slice(0, boundary)); - - if (event) onEvent(event); - - buffer = buffer.slice(boundary + SSE_RECORD_SEPARATOR.length); - boundary = buffer.indexOf(SSE_RECORD_SEPARATOR); - } - } - } - } catch { - // network drop or abort falls through to the reconnect delay - } - - if (signal.aborted) return; - - await new Promise((resolve) => setTimeout(resolve, ModelsService.SSE_RECONNECT_MS)); - } - } - - /** - * Parse one SSE record into its JSON envelope, or null when the record - * carries no data payload or malformed JSON. - */ - private static parseStatusRecord(record: string): ApiModelsSseEvent | null { - const payload = record - .split(SSE_LINE_SEPARATOR) - .filter((line) => line.startsWith(SSE_DATA_PREFIX)) - .map((line) => line.slice(SSE_DATA_PREFIX.length).trim()) - .join(SSE_LINE_SEPARATOR); - - if (payload.length === 0) return null; - - try { - return JSON.parse(payload) as ApiModelsSseEvent; - } catch { - return null; - } - } - - /** - * - * - * Parsing - * - * - */ - /** * Parse a model ID string into its structured components. * @@ -311,4 +202,84 @@ export class ModelsService { return result; } + + /** + * Unload a model (ROUTER mode only). + * Sends POST request to `/models/unload`. Note: the endpoint returns success + * before unloading completes — use polling to await actual unload status. + * + * @param modelId - Model identifier to unload + * @returns Unload response from the server + */ + static async unload(modelId: string): Promise { + return apiPost(API_MODELS.UNLOAD, { model: modelId }); + } + + /** + * Read the /models/sse feed and invoke onEvent for each parsed envelope. + * Reconnects on network drops until the signal aborts. Splits the byte + * stream into SSE records on the blank line boundary; the payload rides in + * the data lines as a JSON envelope with its own model, event and data fields. + */ + static async watchModelEvents( + signal: AbortSignal, + onEvent: (event: ApiModelsSseEvent) => void + ): Promise { + const decoder = new TextDecoder(); + + while (!signal.aborted) { + try { + const response = await fetch(`${base}${API_MODELS.SSE}`, { + headers: getAuthHeaders(), + signal + }); + + if (response.ok && response.body) { + const reader = response.body.getReader(); + + let buffer = ''; + + while (!signal.aborted) { + const { done, value } = await reader.read(); + + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + + const { records, rest } = splitSseRecords(buffer); + + buffer = rest; + + for (const record of records) { + const event = ModelsService.parseStatusRecord(record); + + if (event) onEvent(event); + } + } + } + } catch { + // network drop or abort falls through to the reconnect delay + } + + if (signal.aborted) return; + + await new Promise((resolve) => setTimeout(resolve, ModelsService.SSE_RECONNECT_MS)); + } + } + + /** + * Parse one SSE record into its JSON envelope, or null when the record + * carries no data payload or malformed JSON. + */ + private static parseStatusRecord(record: string): ApiModelsSseEvent | null { + const payload = extractSseDataPayload(record); + + if (payload.length === 0) return null; + + try { + return JSON.parse(payload) as ApiModelsSseEvent; + } catch { + return null; + } + } } diff --git a/tools/ui/src/lib/services/parameter-sync.service.ts b/tools/ui/src/lib/services/parameter-sync.service.ts index 0ed9ebd48ce3..467e7c2dbd7a 100644 --- a/tools/ui/src/lib/services/parameter-sync.service.ts +++ b/tools/ui/src/lib/services/parameter-sync.service.ts @@ -1,3 +1,11 @@ +/** + * ParameterSyncService - Syncs sampling parameters with the server + * + * Decides for each sampling parameter whether the user's setting is an + * override of the server default, and normalizes floating-point values. + * No reactive state; consumed by settingsStore. + */ + import { SETTINGS_KEYS, SYNCABLE_PARAMETERS } from '$lib/constants'; import { ParameterSource, SyncableParameterType } from '$lib/enums'; import type { ParameterInfo, ParameterRecord, ParameterValue } from '$lib/types'; @@ -5,22 +13,47 @@ import { normalizeFloatingPoint } from '$lib/utils'; export class ParameterSyncService { /** + * Check if a parameter can be synced from server. * - * - * Extraction - * - * + * @param key - The parameter key to check + * @returns True if the parameter is in the syncable parameters list */ + static canSyncParameter(key: string): boolean { + return SYNCABLE_PARAMETERS.some((param) => param.key === key && param.canSync); + } /** - * Round floating-point numbers to avoid JavaScript precision issues. - * E.g., 0.1 + 0.2 = 0.30000000000000004 → 0.3 + * Create a diff between current settings and server defaults. + * Shows which parameters differ from server values, useful for debugging + * and for the "Reset to defaults" functionality. * - * @param value - Parameter value to normalize - * @returns Precision-normalized value + * @param currentSettings - Current parameter values in the settings store + * @param serverDefaults - Default values extracted from server props + * @returns Record of parameter diffs with current value, server value, and whether they differ */ - private static roundFloatingPoint(value: ParameterValue): ParameterValue { - return normalizeFloatingPoint(value) as ParameterValue; + static createParameterDiff( + currentSettings: ParameterRecord, + serverDefaults: ParameterRecord + ): Record { + const diff: Record< + string, + { current: ParameterValue; server: ParameterValue; differs: boolean } + > = {}; + + for (const key of this.getSyncableParameterKeys()) { + const currentValue = currentSettings[key]; + const serverValue = serverDefaults[key]; + + if (serverValue !== undefined) { + diff[key] = { + current: currentValue, + differs: currentValue !== serverValue, + server: serverValue + }; + } + } + + return diff; } /** @@ -59,49 +92,6 @@ export class ParameterSyncService { return extracted; } - /** - * - * - * Merging - * - * - */ - - /** - * Merge server defaults with current user settings. - * User overrides always take priority — only parameters not in `userOverrides` - * set will be updated from server defaults. - * - * @param currentSettings - Current parameter values in the settings store - * @param serverDefaults - Default values extracted from server props - * @param userOverrides - Set of parameter keys explicitly overridden by the user - * @returns Merged parameter record with user overrides preserved - */ - static mergeWithServerDefaults( - currentSettings: ParameterRecord, - serverDefaults: ParameterRecord, - userOverrides: Set = new Set() - ): ParameterRecord { - const merged = { ...currentSettings }; - - for (const [key, serverValue] of Object.entries(serverDefaults)) { - // Only update if user hasn't explicitly overridden this parameter - if (!userOverrides.has(key)) { - merged[key] = this.roundFloatingPoint(serverValue); - } - } - - return merged; - } - - /** - * - * - * Info - * - * - */ - /** * Get parameter information including source and values. * Used by SettingsChatParameterSourceIndicator to display the correct badge @@ -133,22 +123,39 @@ export class ParameterSyncService { } /** - * Check if a parameter can be synced from server. + * Get all syncable parameter keys. * - * @param key - The parameter key to check - * @returns True if the parameter is in the syncable parameters list + * @returns Array of parameter keys that can be synced from server */ - static canSyncParameter(key: string): boolean { - return SYNCABLE_PARAMETERS.some((param) => param.key === key && param.canSync); + static getSyncableParameterKeys(): string[] { + return SYNCABLE_PARAMETERS.filter((param) => param.canSync).map((param) => param.key); } /** - * Get all syncable parameter keys. + * Merge server defaults with current user settings. + * User overrides always take priority — only parameters not in `userOverrides` + * set will be updated from server defaults. * - * @returns Array of parameter keys that can be synced from server + * @param currentSettings - Current parameter values in the settings store + * @param serverDefaults - Default values extracted from server props + * @param userOverrides - Set of parameter keys explicitly overridden by the user + * @returns Merged parameter record with user overrides preserved */ - static getSyncableParameterKeys(): string[] { - return SYNCABLE_PARAMETERS.filter((param) => param.canSync).map((param) => param.key); + static mergeWithServerDefaults( + currentSettings: ParameterRecord, + serverDefaults: ParameterRecord, + userOverrides: Set = new Set() + ): ParameterRecord { + const merged = { ...currentSettings }; + + for (const [key, serverValue] of Object.entries(serverDefaults)) { + // Only update if user hasn't explicitly overridden this parameter + if (!userOverrides.has(key)) { + merged[key] = this.roundFloatingPoint(serverValue); + } + } + + return merged; } /** @@ -176,44 +183,13 @@ export class ParameterSyncService { } /** + * Round floating-point numbers to avoid JavaScript precision issues. + * E.g., 0.1 + 0.2 = 0.30000000000000004 → 0.3 * - * - * Diff - * - * - */ - - /** - * Create a diff between current settings and server defaults. - * Shows which parameters differ from server values, useful for debugging - * and for the "Reset to defaults" functionality. - * - * @param currentSettings - Current parameter values in the settings store - * @param serverDefaults - Default values extracted from server props - * @returns Record of parameter diffs with current value, server value, and whether they differ + * @param value - Parameter value to normalize + * @returns Precision-normalized value */ - static createParameterDiff( - currentSettings: ParameterRecord, - serverDefaults: ParameterRecord - ): Record { - const diff: Record< - string, - { current: ParameterValue; server: ParameterValue; differs: boolean } - > = {}; - - for (const key of this.getSyncableParameterKeys()) { - const currentValue = currentSettings[key]; - const serverValue = serverDefaults[key]; - - if (serverValue !== undefined) { - diff[key] = { - current: currentValue, - differs: currentValue !== serverValue, - server: serverValue - }; - } - } - - return diff; + private static roundFloatingPoint(value: ParameterValue): ParameterValue { + return normalizeFloatingPoint(value) as ParameterValue; } } diff --git a/tools/ui/src/lib/services/props.service.ts b/tools/ui/src/lib/services/props.service.ts index 46f4915fadb8..488a67b641cd 100644 --- a/tools/ui/src/lib/services/props.service.ts +++ b/tools/ui/src/lib/services/props.service.ts @@ -1,14 +1,14 @@ +/** + * PropsService - Fetches server properties from /props + * + * Returns global server settings and capabilities, including per-model + * modalities in MODEL mode. No reactive state; consumed by serverStore and + * the model props manager. + */ + import { apiFetchWithParams } from '$lib/utils'; export class PropsService { - /** - * - * - * Fetching - * - * - */ - /** * Fetches global server properties from the `/props` endpoint. * In MODEL mode, returns modalities for the single loaded model. diff --git a/tools/ui/src/lib/services/read-media.service.ts b/tools/ui/src/lib/services/read-media.service.ts index 8de9bbbea597..2858795e8b11 100644 --- a/tools/ui/src/lib/services/read-media.service.ts +++ b/tools/ui/src/lib/services/read-media.service.ts @@ -1,3 +1,10 @@ +/** + * ReadMediaService - Reads local media files for the read_media tool + * + * Encodes image and audio files as base64 data URLs with the metadata the + * model needs. No reactive state; consumed by toolsStore. + */ + import { ToolsService } from './tools.service'; import { FILE_EXTENSION_SEPARATOR, @@ -40,7 +47,7 @@ function fileExtension(path: string): string { * actually use the result - the server has no idea which model is selected. * * @see buildReadMediaToolDefinition in constants/read-media.ts - tool schema sent to the LLM - * @see agenticStore in stores/agentic.svelte.ts - tool dispatch and attachment extraction + * @see agenticStore in stores/agentic/index.svelte.ts - tool dispatch and attachment extraction */ export class ReadMediaService { static async executeTool( diff --git a/tools/ui/src/lib/services/router.service.ts b/tools/ui/src/lib/services/router.service.ts index 59de4cb6fe23..217de38f3a8c 100644 --- a/tools/ui/src/lib/services/router.service.ts +++ b/tools/ui/src/lib/services/router.service.ts @@ -1,3 +1,10 @@ +/** + * RouterService - Builds app route paths + * + * Returns chat and settings route strings from a single source of truth + * (ROUTES). No state. + */ + import { ROUTES } from '$lib/constants'; export class RouterService { diff --git a/tools/ui/src/lib/services/sandbox-harness.ts b/tools/ui/src/lib/services/sandbox-harness.ts index 189ff59a5ec9..29f9ad2a56bf 100644 --- a/tools/ui/src/lib/services/sandbox-harness.ts +++ b/tools/ui/src/lib/services/sandbox-harness.ts @@ -1,3 +1,10 @@ +/** + * Sandbox harness - builds the srcdoc document for the sandboxed iframe + * + * Produces the HTML/CSP/worker shim that runs untrusted model code in an + * opaque origin. Consumed by sandbox.service. + */ + import WORKER_SHIM from './sandbox-worker.js?raw'; import { NEWLINE } from '$lib/constants'; diff --git a/tools/ui/src/lib/services/sandbox.service.ts b/tools/ui/src/lib/services/sandbox.service.ts index 27da9d2634f5..bdc63e4edf18 100644 --- a/tools/ui/src/lib/services/sandbox.service.ts +++ b/tools/ui/src/lib/services/sandbox.service.ts @@ -1,3 +1,11 @@ +/** + * SandboxService - Runs untrusted code in a sandboxed worker + * + * Executes model-generated code inside a CSP-restricted, opaque-origin + * iframe worker with output and timeout limits. No reactive state; consumed + * by toolsStore for code-execution tools. + */ + import { buildSandboxHarness } from './sandbox-harness'; import { NEWLINE, @@ -8,7 +16,7 @@ import { SANDBOX_TOOL_NAME, SANDBOX_TRUNCATION_NOTICE } from '$lib/constants'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { ToolExecutionResult } from '$lib/types'; /** Cached harnesses keyed by whether nerdamer is included. */ diff --git a/tools/ui/src/lib/services/tools.service.ts b/tools/ui/src/lib/services/tools.service.ts index 2b3a2c0dc7a2..78229756ce07 100644 --- a/tools/ui/src/lib/services/tools.service.ts +++ b/tools/ui/src/lib/services/tools.service.ts @@ -1,3 +1,10 @@ +/** + * ToolsService - Stateless server tools API layer + * + * Fetches the server's /tools listing and streams tool execution results. + * No reactive state; consumed by toolsStore. + */ + import { base } from '$app/paths'; import { API_TOOLS, HEADERS } from '$lib/constants'; import { ToolResponseField } from '$lib/enums'; @@ -7,15 +14,6 @@ import { getJsonHeaders } from '$lib/utils/api-headers'; import { parseSseJsonStream, type SseJsonEvent } from '$lib/utils/sse'; export class ToolsService { - /** - * Fetch the list of server tools from the server. - * - * @returns Array of tool definitions in OpenAI-compatible format - */ - static async list(): Promise { - return apiFetch(API_TOOLS.LIST); - } - /** * Execute a server tool on the server. * @@ -76,6 +74,15 @@ export class ToolsService { }); } + /** + * Fetch the list of server tools from the server. + * + * @returns Array of tool definitions in OpenAI-compatible format + */ + static async list(): Promise { + return apiFetch(API_TOOLS.LIST); + } + /** * Stream a server tool's output chunks from the server. The server * `POST /tools` endpoint with `{stream: true}` emits `data: {"chunk": "..."}` diff --git a/tools/ui/src/lib/stores/agentic/gates.svelte.ts b/tools/ui/src/lib/stores/agentic/gates.svelte.ts new file mode 100644 index 000000000000..6b52fa3aff5b --- /dev/null +++ b/tools/ui/src/lib/stores/agentic/gates.svelte.ts @@ -0,0 +1,208 @@ +/** + * AgenticGates - User interaction gates for the agentic loop + * + * Owns the state the loop waits on between turns: tool permission requests, + * turn-limit continue prompts and queued steering messages. The loop awaits + * requestPermission/requestContinue; the UI resolves them through + * resolvePermission/resolveContinue. Owned by agenticStore, no host coupling. + */ + +import { ToolPermissionDecision } from '$lib/enums'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { permissionsStore } from '$lib/stores/permissions.svelte'; +import { toolsStore } from '$lib/stores/tools.svelte'; +import type { DatabaseMessageExtra, SteeringMessage } from '$lib/types'; +import { SvelteMap } from 'svelte/reactivity'; + +export class AgenticGates { + /** Resolve functions for pending continue Promises; nothing derives from this map */ + private continueResolvers = new SvelteMap void>(); + /** Dedicated reactive state for pending continue requests (turn limit reached) */ + private pendingContinueRequests = new SvelteMap(); + + /** Dedicated reactive state for pending permission requests (ensures immediate UI updates) */ + private pendingPermissions = new SvelteMap< + string, + { toolName: string; serverLabel: string } | null + >(); + /** Resolve functions for pending permission Promises; nothing derives from this map */ + private permissionResolvers = new SvelteMap void>(); + + /** Reactive: queued steering messages to inject between turns */ + private steeringMessages = new SvelteMap(); + + /** + * Drop all pending gate state for a conversation, e.g. when a flow exits. + */ + clear(conversationId: string): void { + this.pendingPermissions.set(conversationId, null); + this.permissionResolvers.delete(conversationId); + this.pendingContinueRequests.set(conversationId, false); + this.continueResolvers.delete(conversationId); + this.steeringMessages.delete(conversationId); + } + + /** + * Clear the pending steering message without consuming it. + */ + clearSteeringMessage(conversationId: string): void { + this.steeringMessages.delete(conversationId); + } + + /** + * Consume and return the pending steering message for re-sending. + * Called by chatStore after the agentic flow exits. + */ + consumePendingSteeringMessage(conversationId: string): SteeringMessage | null { + const msg = this.steeringMessages.get(conversationId); + + if (!msg) return null; + + this.steeringMessages.delete(conversationId); + + return msg; + } + + getPendingContinueRequest(conversationId: string): boolean { + return this.pendingContinueRequests.get(conversationId) ?? false; + } + + getPendingPermissionRequest( + conversationId: string + ): { toolName: string; serverLabel: string } | null { + return this.pendingPermissions.get(conversationId) ?? null; + } + + getPendingSteeringMessageContent(conversationId: string): string | null { + return this.steeringMessages.get(conversationId)?.content ?? null; + } + + getPendingSteeringMessageExtras(conversationId: string): DatabaseMessageExtra[] | undefined { + return this.steeringMessages.get(conversationId)?.extras; + } + + hasPendingSteeringMessage(conversationId: string): boolean { + return this.steeringMessages.has(conversationId); + } + + /** + * Queue a steering message. When the current agentic turn completes, + * the flow exits and the caller re-sends the message as a normal chat message. + */ + injectSteeringMessage( + conversationId: string, + content: string, + extras?: DatabaseMessageExtra[] + ): void { + this.steeringMessages.set(conversationId, { content, extras }); + } + + async requestContinue(conversationId: string, signal?: AbortSignal): Promise { + this.pendingContinueRequests.set(conversationId, true); + + return new Promise((resolve) => { + if (signal?.aborted) { + this.pendingContinueRequests.set(conversationId, false); + resolve(false); + + return; + } + + this.continueResolvers.set(conversationId, (shouldContinue) => { + this.pendingContinueRequests.set(conversationId, false); + resolve(shouldContinue); + }); + + signal?.addEventListener( + 'abort', + () => { + const resolver = this.continueResolvers.get(conversationId); + + if (resolver) { + this.continueResolvers.delete(conversationId); + this.pendingContinueRequests.set(conversationId, false); + resolve(false); + } + }, + { once: true } + ); + }); + } + + async requestPermission( + conversationId: string, + toolName: string, + serverLabel: string, + signal?: AbortSignal + ): Promise { + const permissionKey = toolsStore.getPermissionKey(toolName); + + if (permissionKey && permissionsStore.hasTool(permissionKey)) { + return ToolPermissionDecision.ONCE; + } + + this.pendingPermissions.set(conversationId, { serverLabel, toolName }); + + return new Promise((resolve) => { + if (signal?.aborted) { + this.pendingPermissions.set(conversationId, null); + resolve(ToolPermissionDecision.DENY); + + return; + } + + this.permissionResolvers.set(conversationId, (decision) => { + this.pendingPermissions.set(conversationId, null); + + if (decision === ToolPermissionDecision.ALWAYS && permissionKey) { + permissionsStore.allowTool(permissionKey); + } else if (decision === ToolPermissionDecision.ALWAYS_SERVER) { + const serverToolKeys = toolsStore.allTools + .filter((t) => + t.serverName + ? t.serverName === serverLabel + : toolsStore.getToolServerLabel(t.definition.function.name) === serverLabel + ) + .map((t) => toolsStore.getPermissionKey(t.definition.function.name)!) + .filter((k): k is string => k !== null); + + permissionsStore.allowTools(serverToolKeys); + } + + resolve(decision); + }); + + signal?.addEventListener( + 'abort', + () => { + const resolver = this.permissionResolvers.get(conversationId); + + if (resolver) { + this.permissionResolvers.delete(conversationId); + this.pendingPermissions.set(conversationId, null); + resolve(ToolPermissionDecision.DENY); + } + }, + { once: true } + ); + }); + } + + resolveContinue(conversationId: string, shouldContinue: boolean): void { + const resolver = this.continueResolvers.get(conversationId); + + if (resolver) { + this.continueResolvers.delete(conversationId); + resolver(shouldContinue); + } + } + + resolvePermission(conversationId: string, decision: ToolPermissionDecision): void { + const resolver = this.permissionResolvers.get(conversationId); + + if (resolver) { + this.permissionResolvers.delete(conversationId); + resolver(decision); + } + } +} diff --git a/tools/ui/src/lib/stores/agentic.svelte.ts b/tools/ui/src/lib/stores/agentic/index.svelte.ts similarity index 77% rename from tools/ui/src/lib/stores/agentic.svelte.ts rename to tools/ui/src/lib/stores/agentic/index.svelte.ts index d2a2ea88712c..a91e0ba46fa3 100644 --- a/tools/ui/src/lib/stores/agentic.svelte.ts +++ b/tools/ui/src/lib/stores/agentic/index.svelte.ts @@ -1,23 +1,13 @@ /** - * agenticStore - Reactive State Store for Agentic Loop Orchestration + * AgenticStore - Multi-turn agentic loop orchestration * - * Manages multi-turn agentic loop with MCP tools: - * - LLM streaming with tool call detection - * - Tool execution via mcpStore - * - Session state management - * - Turn limit enforcement + * Drives the agentic loop over MCP tools: streams each LLM turn, detects + * tool calls, executes them via mcpStore, and enforces the turn limit. Each + * turn produces one assistant message (with tool_calls) and one tool result + * message per executed call, persisted as separate DB rows. * - * Each agentic turn produces separate DB messages: - * - One assistant message per LLM turn (with tool_calls if any) - * - One tool result message per tool call execution - * - * **Architecture & Relationships:** - * - **ChatService**: Stateless API layer (sendMessage, streaming) - * - **mcpStore**: MCP connection management and tool execution - * - **agenticStore** (this): Reactive state + business logic - * - * @see ChatService in services/chat.service.ts for API operations - * @see mcpStore in stores/mcp.svelte.ts for MCP operations + * Uses ChatService for streaming and mcpStore for tool execution; waits on + * the permission/continue/steering gates owned by {@link AgenticGates}. */ import { DEFAULT_AGENTIC_CONFIG, NEWLINE } from '$lib/constants'; @@ -43,11 +33,11 @@ import { ReadMediaService } from '$lib/services/read-media.service'; import { SandboxService } from '$lib/services/sandbox.service'; import { ToolsService } from '$lib/services/tools.service'; // direct imports between stores, not via the barrel, to avoid circular deps -import { conversationsStore } from '$lib/stores/conversations.svelte'; -import { mcpStore } from '$lib/stores/mcp.svelte'; -import { modelsStore } from '$lib/stores/models.svelte'; -import { permissionsStore } from '$lib/stores/permissions.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { AgenticGates } from '$lib/stores/agentic/gates.svelte'; +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import { mcpStore } from '$lib/stores/mcp/index.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { toolsStore } from '$lib/stores/tools.svelte'; import type { AgenticConfig, @@ -152,141 +142,140 @@ function toAgenticMessages(messages: ApiChatMessageData[]): AgenticMessage[] { } class AgenticStore { - private _sessions = new SvelteMap(); - /** Dedicated reactive state for pending permission requests (ensures immediate UI updates) */ - private _pendingPermissions = new SvelteMap< - string, - { toolName: string; serverLabel: string } | null - >(); - /** Non-reactive: stores resolve functions for pending permission Promises */ - private _permissionResolvers = new Map void>(); - - /** Dedicated reactive state for pending continue requests (turn limit reached) */ - private _pendingContinueRequests = new SvelteMap(); - /** Non-reactive: stores resolve functions for pending continue Promises */ - private _continueResolvers = new Map void>(); - - /** Reactive: queued steering messages to inject between turns */ - private _steeringMessages = new SvelteMap(); + // permission, continue and steering gates the loop waits on between turns + private gates = new AgenticGates(); + private sessions = new SvelteMap(); - get isReady(): boolean { - return true; - } get isAnyRunning(): boolean { - for (const session of this._sessions.values()) { + for (const session of this.sessions.values()) { if (session.isRunning) return true; } return false; } - getSession(conversationId: string): AgenticSession { - let session = this._sessions.get(conversationId); + get isReady(): boolean { + return true; + } - if (!session) { - session = createDefaultSession(); - this._sessions.set(conversationId, session); - } + clearError(conversationId: string): void { + this.updateSession(conversationId, { lastError: null }); + } - return session; + clearSession(conversationId: string): void { + this.sessions.delete(conversationId); } - private updateSession(conversationId: string, update: Partial): void { - const session = this.getSession(conversationId); + /** + * Clear the pending steering message without consuming it. + */ + clearSteeringMessage(conversationId: string): void { + this.gates.clearSteeringMessage(conversationId); + } - this._sessions.set(conversationId, { ...session, ...update }); + constructor() { + // drop per-conversation session state when the conversation is deleted, + // otherwise every conversation that ever ran a flow leaks a session here + conversationsStore.onConversationsDeleted((convIds) => { + for (const convId of convIds) { + this.sessions.delete(convId); + } + }); } - clearSession(conversationId: string): void { - this._sessions.delete(conversationId); + /** + * Consume and return the pending steering message for re-sending. + * Called by chatStore after the agentic flow exits. + */ + consumePendingSteeringMessage(conversationId: string): SteeringMessage | null { + return this.gates.consumePendingSteeringMessage(conversationId); } getActiveSessions(): Array<{ conversationId: string; session: AgenticSession }> { const active: Array<{ conversationId: string; session: AgenticSession }> = []; - for (const [conversationId, session] of this._sessions.entries()) { + for (const [conversationId, session] of this.sessions.entries()) { if (session.isRunning) active.push({ conversationId, session }); } return active; } - isRunning(conversationId: string): boolean { - return this._sessions.get(conversationId)?.isRunning ?? false; - } + getConfig(settings: SettingsConfigType, perChatOverrides?: McpServerOverride[]): AgenticConfig { + const maxTurns = Number(settings.agenticMaxTurns) || DEFAULT_AGENTIC_CONFIG.maxTurns; + const hasTools = + mcpStore.hasEnabledServers(perChatOverrides) || + toolsStore.serverTools.length > 0 || + toolsStore.browserTools.length > 0 || + toolsStore.customTools.length > 0; - // read-only: safe to call from derivations, unlike getSession - getLiveLlmTotals(conversationId: string): AgenticSession['liveLlm'] { - return this._sessions.get(conversationId)?.liveLlm ?? null; + return { + enabled: hasTools && DEFAULT_AGENTIC_CONFIG.enabled, + maxTurns + }; } - // read-only: safe to call from derivations, unlike getSession - getFlowRootMessageId(conversationId: string): string | null { - return this._sessions.get(conversationId)?.flowRootMessageId ?? null; + getCurrentTurn(conversationId: string): number { + return this.sessions.get(conversationId)?.currentTurn ?? 0; } - currentTurn(conversationId: string): number { - return this._sessions.get(conversationId)?.currentTurn ?? 0; + getExecutingToolCallId(conversationId: string): string | null { + return this.sessions.get(conversationId)?.executingToolCallId ?? null; } - totalToolCalls(conversationId: string): number { - return this._sessions.get(conversationId)?.totalToolCalls ?? 0; + // read-only: safe to call from derivations, unlike getSession + getFlowRootMessageId(conversationId: string): string | null { + return this.sessions.get(conversationId)?.flowRootMessageId ?? null; } - lastError(conversationId: string): Error | null { - return this._sessions.get(conversationId)?.lastError ?? null; + getLastError(conversationId: string): Error | null { + return this.sessions.get(conversationId)?.lastError ?? null; } - streamingToolCall(conversationId: string): { name: string; arguments: string } | null { - return this._sessions.get(conversationId)?.streamingToolCall ?? null; + // read-only: safe to call from derivations, unlike getSession + getLiveLlmTotals(conversationId: string): AgenticSession['liveLlm'] { + return this.sessions.get(conversationId)?.liveLlm ?? null; } - executingToolCallId(conversationId: string): string | null { - return this._sessions.get(conversationId)?.executingToolCallId ?? null; + getPendingContinueRequest(conversationId: string): boolean { + return this.gates.getPendingContinueRequest(conversationId); } - pendingPermissionRequest( + getPendingPermissionRequest( conversationId: string ): { toolName: string; serverLabel: string } | null { - return this._pendingPermissions.get(conversationId) ?? null; + return this.gates.getPendingPermissionRequest(conversationId); } - pendingContinueRequest(conversationId: string): boolean { - return this._pendingContinueRequests.get(conversationId) ?? false; + getPendingSteeringMessageContent(conversationId: string): string | null { + return this.gates.getPendingSteeringMessageContent(conversationId); } - resolveContinue(conversationId: string, shouldContinue: boolean): void { - const resolver = this._continueResolvers.get(conversationId); - - if (resolver) { - this._continueResolvers.delete(conversationId); - resolver(shouldContinue); - } + getPendingSteeringMessageExtras(conversationId: string): DatabaseMessageExtra[] | undefined { + return this.gates.getPendingSteeringMessageExtras(conversationId); } - resolvePermission(conversationId: string, decision: ToolPermissionDecision): void { - const resolver = this._permissionResolvers.get(conversationId); + getSession(conversationId: string): AgenticSession { + let session = this.sessions.get(conversationId); - if (resolver) { - this._permissionResolvers.delete(conversationId); - resolver(decision); + if (!session) { + session = createDefaultSession(); + this.sessions.set(conversationId, session); } - } - clearError(conversationId: string): void { - this.updateSession(conversationId, { lastError: null }); + return session; } - hasPendingSteeringMessage(conversationId: string): boolean { - return this._steeringMessages.has(conversationId); + getStreamingToolCall(conversationId: string): { name: string; arguments: string } | null { + return this.sessions.get(conversationId)?.streamingToolCall ?? null; } - pendingSteeringMessageContent(conversationId: string): string | null { - return this._steeringMessages.get(conversationId)?.content ?? null; + getTotalToolCalls(conversationId: string): number { + return this.sessions.get(conversationId)?.totalToolCalls ?? 0; } - pendingSteeringMessageExtras(conversationId: string): DatabaseMessageExtra[] | undefined { - return this._steeringMessages.get(conversationId)?.extras; + hasPendingSteeringMessage(conversationId: string): boolean { + return this.gates.hasPendingSteeringMessage(conversationId); } /** @@ -298,143 +287,19 @@ class AgenticStore { content: string, extras?: DatabaseMessageExtra[] ): void { - this._steeringMessages.set(conversationId, { content, extras }); - } - - /** - * Clear the pending steering message without consuming it. - */ - clearSteeringMessage(conversationId: string): void { - this._steeringMessages.delete(conversationId); - } - - /** - * Consume and return the pending steering message for re-sending. - * Called by chatStore after the agentic flow exits. - */ - consumePendingSteeringMessage(conversationId: string): SteeringMessage | null { - const msg = this._steeringMessages.get(conversationId); - - if (!msg) return null; - - this._steeringMessages.delete(conversationId); - - return msg; + this.gates.injectSteeringMessage(conversationId, content, extras); } - getConfig(settings: SettingsConfigType, perChatOverrides?: McpServerOverride[]): AgenticConfig { - const maxTurns = Number(settings.agenticMaxTurns) || DEFAULT_AGENTIC_CONFIG.maxTurns; - const hasTools = - mcpStore.hasEnabledServers(perChatOverrides) || - toolsStore.serverTools.length > 0 || - toolsStore.browserTools.length > 0 || - toolsStore.customTools.length > 0; - - return { - enabled: hasTools && DEFAULT_AGENTIC_CONFIG.enabled, - maxTurns - }; - } - - private parseToolArguments(args: string | Record): Record { - if (typeof args === 'object') return args; - - const trimmed = args.trim(); - - if (trimmed === '') return {}; - - return JSON.parse(trimmed) as Record; + isRunning(conversationId: string): boolean { + return this.sessions.get(conversationId)?.isRunning ?? false; } - private async requestPermission( - conversationId: string, - toolName: string, - serverLabel: string, - signal?: AbortSignal - ): Promise { - const permissionKey = toolsStore.getPermissionKey(toolName); - - if (permissionKey && permissionsStore.hasTool(permissionKey)) { - return ToolPermissionDecision.ONCE; - } - - this._pendingPermissions.set(conversationId, { serverLabel, toolName }); - - return new Promise((resolve) => { - if (signal?.aborted) { - this._pendingPermissions.set(conversationId, null); - resolve(ToolPermissionDecision.DENY); - - return; - } - - this._permissionResolvers.set(conversationId, (decision) => { - this._pendingPermissions.set(conversationId, null); - - if (decision === ToolPermissionDecision.ALWAYS && permissionKey) { - permissionsStore.allowTool(permissionKey); - } else if (decision === ToolPermissionDecision.ALWAYS_SERVER) { - const serverToolKeys = toolsStore.allTools - .filter((t) => - t.serverName - ? t.serverName === serverLabel - : toolsStore.getToolServerLabel(t.definition.function.name) === serverLabel - ) - .map((t) => toolsStore.getPermissionKey(t.definition.function.name)!) - .filter((k): k is string => k !== null); - - permissionsStore.allowTools(serverToolKeys); - } - - resolve(decision); - }); - - signal?.addEventListener( - 'abort', - () => { - const resolver = this._permissionResolvers.get(conversationId); - - if (resolver) { - this._permissionResolvers.delete(conversationId); - this._pendingPermissions.set(conversationId, null); - resolve(ToolPermissionDecision.DENY); - } - }, - { once: true } - ); - }); + resolveContinue(conversationId: string, shouldContinue: boolean): void { + this.gates.resolveContinue(conversationId, shouldContinue); } - private async requestContinue(conversationId: string, signal?: AbortSignal): Promise { - this._pendingContinueRequests.set(conversationId, true); - - return new Promise((resolve) => { - if (signal?.aborted) { - this._pendingContinueRequests.set(conversationId, false); - resolve(false); - - return; - } - - this._continueResolvers.set(conversationId, (shouldContinue) => { - this._pendingContinueRequests.set(conversationId, false); - resolve(shouldContinue); - }); - - signal?.addEventListener( - 'abort', - () => { - const resolver = this._continueResolvers.get(conversationId); - - if (resolver) { - this._continueResolvers.delete(conversationId); - this._pendingContinueRequests.set(conversationId, false); - resolve(false); - } - }, - { once: true } - ); - }); + resolvePermission(conversationId: string, decision: ToolPermissionDecision): void { + this.gates.resolvePermission(conversationId, decision); } async runAgenticFlow(params: AgenticFlowParams): Promise { @@ -449,11 +314,7 @@ class AgenticStore { } = params; // Clear any pending permissions/continue requests for this conversation when starting a new flow - this._pendingPermissions.set(conversationId, null); - this._permissionResolvers.delete(conversationId); - this._pendingContinueRequests.set(conversationId, false); - this._continueResolvers.delete(conversationId); - this._steeringMessages.delete(conversationId); + this.gates.clear(conversationId); // Ensure server tools are fetched before checking if agentic is enabled if (toolsStore.serverTools.length === 0 && !toolsStore.loading) { @@ -482,26 +343,8 @@ class AgenticStore { console.log(`[AgenticStore] Starting agentic flow with ${tools.length} tools`); - const normalizedMessages: ApiChatMessageData[] = ( - await Promise.all( - messages.map((msg) => { - if ('id' in msg && 'convId' in msg && 'timestamp' in msg) - return ChatService.convertDbMessageToApiChatMessageData( - msg as DatabaseMessage & { extra?: DatabaseMessageExtra[] } - ); - - return msg as ApiChatMessageData; - }) - ) - ).filter((msg: { role: ChatRole; content: string | ApiChatMessageContentPart[] }) => { - if (msg.role === MessageRole.SYSTEM) { - const content = typeof msg.content === 'string' ? msg.content : ''; - - return content.trim().length > 0; - } - - return true; - }); + const normalizedMessages: ApiChatMessageData[] = + await ChatService.normalizeMessagesForApi(messages); this.updateSession(conversationId, { currentTurn: 0, @@ -550,6 +393,30 @@ class AgenticStore { } } + private buildAttachmentName(mimeType: string, index: number): string { + const extension = mimeType.startsWith(MimeTypePrefix.AUDIO) + ? (AUDIO_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_AUDIO_EXTENSION) + : (IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION); + + return `${MCP_ATTACHMENT_NAME_PREFIX}-${Date.now()}-${index}.${extension}`; + } + + private buildFinalTimings( + capturedTimings: ChatMessageTimings | undefined, + agenticTimings: ChatMessageAgenticTimings + ): ChatMessageTimings | undefined { + if (agenticTimings.toolCallsCount === 0) return capturedTimings; + + return { + agentic: agenticTimings, + cache_n: capturedTimings?.cache_n, + predicted_ms: capturedTimings?.predicted_ms, + predicted_n: capturedTimings?.predicted_n, + prompt_ms: capturedTimings?.prompt_ms, + prompt_n: capturedTimings?.prompt_n + }; + } + private async executeAgenticLoop(params: { conversationId: string; messages: ApiChatMessageData[]; @@ -596,7 +463,7 @@ class AgenticStore { while (true) { if (turn >= maxTurns) { // Turn limit reached - ask user whether to continue - const shouldContinue = await this.requestContinue(conversationId, signal); + const shouldContinue = await this.gates.requestContinue(conversationId, signal); // Yield to allow Svelte to flush the UI update await new Promise((r) => setTimeout(r, 0)); @@ -769,7 +636,7 @@ class AgenticStore { // === Steering check: if a user message was queued during this turn, exit the flow. // The caller (chatStore) will consume the pending message and re-send it normally. - if (this._steeringMessages.has(conversationId)) { + if (this.gates.hasPendingSteeringMessage(conversationId)) { console.log('[AgenticStore] Steering message detected after turn, exiting agentic flow'); await onAssistantTurnComplete?.( turnContent, @@ -847,7 +714,7 @@ class AgenticStore { } // Check for pending steering message - skip remaining tool calls - if (this._steeringMessages.has(conversationId)) { + if (this.gates.hasPendingSteeringMessage(conversationId)) { console.log( `[AgenticStore] Steering message detected, skipping ${normalizedCalls.length - i} remaining tool call(s)` ); @@ -872,7 +739,7 @@ class AgenticStore { const toolName = toolCall.function.name; const serverLabel = toolsStore.getToolServerLabel(toolName); // Ask for permission before executing the tool - const permission = await this.requestPermission( + const permission = await this.gates.requestPermission( conversationId, toolName, serverLabel, @@ -959,8 +826,8 @@ class AgenticStore { executionResult = await ReadMediaService.executeTool( args, { - audio: modelsStore.modelSupportsAudio(effectiveModel), - vision: modelsStore.modelSupportsVision(effectiveModel) + audio: modelsStore.props.modelSupportsAudio(effectiveModel), + vision: modelsStore.props.modelSupportsVision(effectiveModel) }, signal, conversationsStore.activeConversation?.cwd @@ -1058,7 +925,7 @@ class AgenticStore { for (const attachment of attachments) { if (attachment.type === AttachmentType.AUDIO) { - if (modelsStore.modelSupportsAudio(effectiveModel)) { + if (modelsStore.props.modelSupportsAudio(effectiveModel)) { contentParts.push({ input_audio: { data: (attachment as DatabaseMessageExtraAudioFile).base64Data, @@ -1070,7 +937,7 @@ class AgenticStore { }); } } else if (attachment.type === AttachmentType.IMAGE) { - if (modelsStore.modelSupportsVision(effectiveModel)) { + if (modelsStore.props.modelSupportsVision(effectiveModel)) { contentParts.push({ image_url: { url: (attachment as DatabaseMessageExtraImageFile).base64Url @@ -1101,7 +968,7 @@ class AgenticStore { } // If tools were interrupted by a steering message, exit now instead of starting another LLM turn - if (this._steeringMessages.has(conversationId)) { + if (this.gates.hasPendingSteeringMessage(conversationId)) { console.log( '[AgenticStore] Steering message detected after tool execution, exiting agentic flow' ); @@ -1114,35 +981,6 @@ class AgenticStore { } } - private buildFinalTimings( - capturedTimings: ChatMessageTimings | undefined, - agenticTimings: ChatMessageAgenticTimings - ): ChatMessageTimings | undefined { - if (agenticTimings.toolCallsCount === 0) return capturedTimings; - - return { - agentic: agenticTimings, - cache_n: capturedTimings?.cache_n, - predicted_ms: capturedTimings?.predicted_ms, - predicted_n: capturedTimings?.predicted_n, - prompt_ms: capturedTimings?.prompt_ms, - prompt_n: capturedTimings?.prompt_n - }; - } - - private normalizeToolCalls(toolCalls: ApiChatCompletionToolCall[]): AgenticToolCallList { - if (!toolCalls) return []; - - return toolCalls.map((call, index) => ({ - function: { - arguments: call?.function?.arguments ?? '', - name: call?.function?.name ?? '' - }, - id: call?.id ?? `tool_${index}`, - type: (call?.type as ToolCallType.FUNCTION) ?? ToolCallType.FUNCTION - })); - } - private extractBase64Attachments(result: string): { cleanedResult: string; attachments: DatabaseMessageExtra[]; @@ -1198,12 +1036,33 @@ class AgenticStore { return { attachments, cleanedResult: cleanedLines.join(NEWLINE) }; } - private buildAttachmentName(mimeType: string, index: number): string { - const extension = mimeType.startsWith(MimeTypePrefix.AUDIO) - ? (AUDIO_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_AUDIO_EXTENSION) - : (IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION); + private normalizeToolCalls(toolCalls: ApiChatCompletionToolCall[]): AgenticToolCallList { + if (!toolCalls) return []; - return `${MCP_ATTACHMENT_NAME_PREFIX}-${Date.now()}-${index}.${extension}`; + return toolCalls.map((call, index) => ({ + function: { + arguments: call?.function?.arguments ?? '', + name: call?.function?.name ?? '' + }, + id: call?.id ?? `tool_${index}`, + type: (call?.type as ToolCallType.FUNCTION) ?? ToolCallType.FUNCTION + })); + } + + private parseToolArguments(args: string | Record): Record { + if (typeof args === 'object') return args; + + const trimmed = args.trim(); + + if (trimmed === '') return {}; + + return JSON.parse(trimmed) as Record; + } + + private updateSession(conversationId: string, update: Partial): void { + const session = this.getSession(conversationId); + + this.sessions.set(conversationId, { ...session, ...update }); } } diff --git a/tools/ui/src/lib/stores/chat.svelte.ts b/tools/ui/src/lib/stores/chat.svelte.ts deleted file mode 100644 index b7add77779f4..000000000000 --- a/tools/ui/src/lib/stores/chat.svelte.ts +++ /dev/null @@ -1,2868 +0,0 @@ -/** - * chatStore - Reactive State Store for Chat Operations - * - * Manages chat lifecycle, streaming, message operations, and processing state. - * - * **Architecture & Relationships:** - * - **ChatService**: Stateless API layer (sendMessage, streaming) - * - **chatStore** (this): Reactive state + business logic - * - **conversationsStore**: Conversation persistence and navigation - * - * @see ChatService in services/chat.service.ts for API operations - */ - -import { - CONVERSATION_ID_SEPARATOR, - CWD_CLEARED_TEXT, - INACTIVE_CONVERSATION, - STREAM_RESUME_RETRY_MS, - SYSTEM_MESSAGE_PLACEHOLDER, - TITLE_GENERATION -} from '$lib/constants'; -import { - ContinueIntentKind, - ErrorDialogType, - MessageRole, - MessageType, - ReasoningEffort, - StreamConnectionState -} from '$lib/enums'; -import { ChatService } from '$lib/services/chat.service'; -import { DatabaseService } from '$lib/services/database.service'; -// direct imports between stores, not via the barrel, to avoid circular deps -import { agenticStore } from '$lib/stores/agentic.svelte'; -import { conversationsStore } from '$lib/stores/conversations.svelte'; -import { mcpStore } from '$lib/stores/mcp.svelte'; -import { modelsStore } from '$lib/stores/models.svelte'; -import { serverStore } from '$lib/stores/server.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; -import { toolsStore } from '$lib/stores/tools.svelte'; -import type { - ApiChatMessageData, - ApiProcessingState, - ApiStreamSession, - ChatMessagePromptProgress, - ChatMessageTimings, - ChatStreamCallbacks, - DatabaseMessage, - DatabaseMessageExtra, - ErrorDialogState -} from '$lib/types'; -import { - classifyContinueIntent, - filterByLeafNodeId, - findDescendantMessages, - findLeafNode, - findMessageById, - formatCwdMessage, - generateConversationTitle, - getConversationModel, - isAbortError, - normalizeModelName, - streamIdentity -} from '$lib/utils'; -import { SvelteMap, SvelteSet } from 'svelte/reactivity'; - -interface ConversationStateEntry { - lastAccessed: number; -} - -class ChatStore { - activeProcessingState = $state(null); - currentResponse = $state(''); - errorDialogState = $state(null); - isLoading = $state(false); - // true while the active conversation streams reasoning content but no visible content yet - isReasoning = $state(false); - // resumable stream connection state for the active conversation - // streaming -> bytes flowing normally, resuming -> waiting on /v1/stream reconnect, lost -> unrecoverable - streamConnectionState = $state(StreamConnectionState.STREAMING); - chatLoadingStates = new SvelteMap(); - chatReasoningStates = new SvelteMap(); - chatStreamingStates = new SvelteMap< - string, - { response: string; messageId: string; model?: string | null } - >(); - // convs that the backend reports as having a running session, populated by the global sync - // at app mount and on visibilitychange. it does not overlap with chatLoadingStates which - // tracks inferences driven by this browser, both are unioned to feed the sidebar spinners - private remoteRunningConvs = new SvelteSet(); - // per conv attach lifecycle, used to derive the global streaming flag without flipping it - // off when one conv finishes while another is still streaming. mirrors chatLoadingStates - // in scope but tracks the attach + tee replay path specifically - private attachingConvs = new SvelteSet(); - // pending resume retry timers while an owning model loads, one per conv - private resumeRetryTimers = new SvelteMap>(); - // convs whose resume waits on a model load: their loading state belongs to the retry loop, - // so discoverActiveStream must not treat it as a live send and bail - private resumePendingConvs = new SvelteSet(); - // in-flight discoverActiveStream guard, keyed by conv id - private discoveringConvs = new SvelteSet(); - private abortControllers = new SvelteMap(); - private preEncodeAbortController: AbortController | null = null; - private processingStates = new SvelteMap(); - private conversationStateTimestamps = new SvelteMap(); - private activeConversationId = $state(null); - private isStreamingActive = $state(false); - private isEditModeActive = $state(false); - private addFilesHandler: ((files: File[]) => void) | null = $state(null); - pendingEditMessageId = $state(null); - private _pendingDraftMessage = $state(''); - private _pendingDraftFiles = $state([]); - - /** Reactive: queued pending messages for non-agentic streaming */ - private _pendingMessages = new SvelteMap< - string, - { content: string; extras?: DatabaseMessageExtra[] } - >(); - - private setChatLoading(convId: string, loading: boolean): void { - this.touchConversationState(convId); - - if (loading) { - this.chatLoadingStates.set(convId, true); - - if (convId === conversationsStore.activeConversation?.id) this.isLoading = true; - } else { - this.chatLoadingStates.delete(convId); - - if (convId === conversationsStore.activeConversation?.id) this.isLoading = false; - - this.setChatReasoning(convId, false); - // the local pipe is the authoritative observer of session end: when it finishes (clean - // onComplete or explicit Stop), the backend session is finalized too, so we drop the - // sidebar hint for this conv right away instead of waiting for the next visibilitychange - // snapshot. without this the spinner ghosts until the user toggles the tab - this.remoteRunningConvs.delete(convId); - } - } - - private setChatReasoning(convId: string, reasoning: boolean): void { - if (reasoning) { - this.chatReasoningStates.set(convId, true); - - if (convId === conversationsStore.activeConversation?.id) this.isReasoning = true; - } else { - this.chatReasoningStates.delete(convId); - - if (convId === conversationsStore.activeConversation?.id) this.isReasoning = false; - } - } - private setChatStreaming( - convId: string, - response: string, - messageId: string, - model?: string | null - ): void { - this.touchConversationState(convId); - this.chatStreamingStates.set(convId, { - messageId, - model: model ?? this.chatStreamingStates.get(convId)?.model, - response - }); - - if (convId === conversationsStore.activeConversation?.id) this.currentResponse = response; - } - private clearChatStreaming(convId: string, messageId?: string): void { - // session aware: a stale generation must not wipe a newer one's streaming state on the - // same conversation, that would drop the frozen stop identity and stop the wrong session - if (messageId !== undefined) { - const cur = this.chatStreamingStates.get(convId); - - if (cur && cur.messageId !== messageId) return; - } - - this.chatStreamingStates.delete(convId); - - if (convId === conversationsStore.activeConversation?.id) this.currentResponse = ''; - } - private getChatStreamingState( - convId: string - ): { response: string; messageId: string } | undefined { - return this.chatStreamingStates.get(convId); - } - syncLoadingStateForChat(convId: string): void { - this.isLoading = this.chatLoadingStates.get(convId) || false; - this.isReasoning = this.chatReasoningStates.get(convId) || false; - const s = this.chatStreamingStates.get(convId); - - this.currentResponse = s?.response || ''; - this.isStreamingActive = s !== undefined; - this.setActiveProcessingConversation(convId); - - // Sync streaming content to activeMessages so UI displays current content - if (s?.response && s?.messageId) { - const idx = conversationsStore.findMessageIndex(s.messageId); - - if (idx !== -1) { - conversationsStore.updateMessageAtIndex(idx, { content: s.response }); - } - } - } - /** - * Server side stream discovery, split in three pieces: - * - * probeServerStream(convId) -> hits POST /v1/streams/lookup with the conv id, returns the session to attach - * to or null. Pure read, no side effect, no UI lock. Safe to fire in parallel with anything. - * - * attachServerStream(convId) -> flips the spinner immediately, fetches the replay stream - * from byte 0, finds the assistant slot to splice into (creates a placeholder if the conv has - * no assistant message yet, for cross device or fresh local DB cases), and pipes the SSE bytes - * into the message via handleStreamResponse. - * - * discoverActiveStream(convId) -> probe + attach in one call. Used by callers that do not need - * to overlap the probe with other async work. - * - * The mount of the chat page in +page.svelte calls probeServerStream in parallel with - * loadConversation, then attachServerStream once both have settled. This gives the earliest - * possible time to spinner and avoids racing against an empty activeMessages array. - */ - async probeServerStream(convId: string): Promise { - if (!convId) return null; - - let sessions: ApiStreamSession[]; - - try { - sessions = await ChatService.lookupStreamSessions([convId]); - } catch (e) { - console.warn(`probeServerStream failed for conv ${convId}:`, e); - - return null; - } - - return ChatService.selectActiveStream(sessions); - } - - async attachServerStream(convId: string, streamId?: string): Promise { - if (!convId) return; - - if (this.chatStreamingStates.has(convId)) return; - - // flip the spinner immediately, the user sees activity as soon as the conv becomes active. - // the global isStreamingActive flag is derived from attachingConvs.size, so adding here - // turns it on, and removing in unlock only turns it off when this is the last attach - this.setChatLoading(convId, true); - this.attachingConvs.add(convId); - this.setStreamingActive(true); - - // only set the active processing conv if we are looking at it, otherwise a background - // attach would steal the indicator from the conv the user is currently viewing - if (convId === conversationsStore.activeConversation?.id) { - this.setActiveProcessingConversation(convId); - } - - const unlock = () => { - this.attachingConvs.delete(convId); - - // flip the global flag off only when no other conv is still attaching - if (this.attachingConvs.size === 0) { - this.setStreamingActive(false); - } - - this.setChatLoading(convId, false); - this.clearChatStreaming(convId); - }; - // fetch the replay stream from byte 0, rebuild the assistant message from scratch. - // resolve the server side identity, fall back to streamIdentity when the caller does not - // pass a streamId. probeServerStream returns the full id (with ::model suffix when present) - const id = streamId || streamIdentity(convId, modelsStore.selectedModelName); - - let response: Response; - - try { - response = await ChatService.fetchStreamReplay(id); - } catch (e) { - console.error(`attachServerStream replay failed for conv ${convId}:`, e); - unlock(); - - return; - } - - // load the target conversation messages by id, not via the active store. when multiple - // attaches run in parallel the active store may reflect another conv and writing through - // its index mixes content across convs (CoT flicker, message bleed). by going through the - // DB we stay isolated, and only mirror into the active store when the attached conv is - // the one currently displayed - let messages: DatabaseMessage[]; - - try { - messages = await DatabaseService.getConversationMessages(convId); - } catch (e) { - console.error('attachServerStream load messages failed:', e); - unlock(); - - return; - } - - // locate the slot to splice into, create a placeholder assistant message if there is none. - // we use the conv-scoped findLastAssistantIdx helpers, they only depend on the array - let targetIdx = this.findLastAssistantIdx(messages); - - if (targetIdx === -1) { - const lastUserIdx = this.findLastUserIdx(messages); - - if (lastUserIdx === -1) { - console.warn( - `attachServerStream: conv ${convId} has no user or assistant message, cannot splice` - ); - unlock(); - - return; - } - - try { - const placeholder = await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId, - parent: messages[lastUserIdx].id, - role: MessageRole.ASSISTANT, - timestamp: Date.now(), - toolCalls: '', - type: MessageType.TEXT - } as Omit, - messages[lastUserIdx].id - ); - - messages = [...messages, placeholder]; - targetIdx = messages.length - 1; - - // only push into the active store when this conv is the one displayed right now - if (convId === conversationsStore.activeConversation?.id) { - conversationsStore.addMessageToActive(placeholder); - } - } catch (e) { - console.error('attachServerStream placeholder creation failed:', e); - unlock(); - - return; - } - } - - if (targetIdx === -1) { - unlock(); - - return; - } - - const targetMessage = messages[targetIdx]; - const targetMessageId = targetMessage.id; - // when the assistant slot already has content, the running session is a continue or - // another append flow and its buffer holds only the appended deltas. preserve the prefix - // and let the replay add to it. when the slot is empty the session buffer holds the whole - // message so we wipe and rebuild from byte 0 - const existingContent = targetMessage.content ?? ''; - const existingReasoning = targetMessage.reasoningContent ?? ''; - const isAppendMode = existingContent.length > 0; - // helper: write to the active store only when the attached conv is currently displayed. - // the lookup by message id is robust to reordering of activeMessages, two parallel attaches - // can no longer step on each other's indices - const writeActive = (updates: Partial) => { - if (convId !== conversationsStore.activeConversation?.id) { - return; - } - - const liveIdx = conversationsStore.findMessageIndex(targetMessageId); - - if (liveIdx === -1) return; - - conversationsStore.updateMessageAtIndex(liveIdx, updates); - }; - - if (!isAppendMode) { - writeActive({ content: '', reasoningContent: undefined }); - } - - // extract the model suffix, the resume calls in handleStreamResponse must reuse the model - // the session was tagged with, not the live dropdown - const sepIdx = id.indexOf(CONVERSATION_ID_SEPARATOR); - const attachedModel: string | null = sepIdx === -1 ? null : id.slice(sepIdx + 2); - - this.setChatStreaming(convId, existingContent, targetMessageId, attachedModel); - const abortController = this.getOrCreateAbortController(convId); - - let streamedContent = ''; - let streamedReasoningContent = ''; - - const cleanup = () => { - unlock(); - this.setProcessingState(convId, null); - }; - - try { - await ChatService.handleStreamResponse( - response, - (chunk: string) => { - streamedContent += chunk; - const displayed = isAppendMode ? existingContent + streamedContent : streamedContent; - - writeActive({ content: displayed }); - this.setChatStreaming(convId, displayed, targetMessageId); - }, - async ( - finalContent?: string, - reasoningContent?: string, - timings?: ChatMessageTimings, - toolCalls?: string - ) => { - const streamed = streamedContent || finalContent || ''; - const streamedR = streamedReasoningContent || reasoningContent || ''; - const content = isAppendMode ? existingContent + streamed : streamed; - const reasoning = isAppendMode ? existingReasoning + streamedR : streamedR; - - // the DB write is the source of truth, mirror to the active store only when - // the conv is currently displayed - await DatabaseService.updateMessage(targetMessageId, { - content, - reasoningContent: reasoning || undefined, - timings, - toolCalls: toolCalls || '' - }); - writeActive({ - content, - reasoningContent: reasoning || undefined, - timings - }); - cleanup(); - }, - (err: Error) => { - console.error('attachServerStream pipe error:', err); - cleanup(); - }, - (chunk: string) => { - streamedReasoningContent += chunk; - const displayed = isAppendMode - ? existingReasoning + streamedReasoningContent - : streamedReasoningContent; - - writeActive({ reasoningContent: displayed }); - }, - undefined, - undefined, - undefined, - undefined, - convId, - abortController.signal, - (connState: StreamConnectionState) => { - if (convId === conversationsStore.activeConversation?.id) { - this.streamConnectionState = connState; - } - }, - attachedModel - ); - } catch (e) { - console.error('attachServerStream pipe crashed:', e); - cleanup(); - } - } - - /** - * Model frozen at send time for a stream awaiting resume, from the persisted stream state. - * The load progress indicator targets it after a reload, when the message row has no model - * yet and the dropdown selection may not be restored. - */ - getResumeModel(convId: string): string | null { - return ChatService.getStreamState(convId)?.model ?? null; - } - - async discoverActiveStream(convId: string): Promise { - if (!convId) return; - - if (this.chatStreamingStates.has(convId)) return; - - if (this.chatLoadingStates.get(convId) && !this.resumePendingConvs.has(convId)) return; - - // concurrency guard: another discover may already be running for this conv (typical race - // between mount and visibilitychange on tab switch). a second concurrent fetch on the same - // /v1/stream would duplicate every byte into the DB message, this guard bounces it - if (this.discoveringConvs.has(convId)) return; - - this.discoveringConvs.add(convId); - - try { - // the model is frozen at POST time, rebuild the exact conv::model identity from the - // persisted state so the lookup key matches what the server stored. null means a single - // model conv with no ::suffix, only guess from the dropdown with no persisted state - const localState = ChatService.getStreamState(convId); - const streamId = ChatService.resumeStreamIdentity( - convId, - localState, - modelsStore.selectedModelName - ); - // primary path: ask the server which sessions exist for this identity - const serverTarget = await this.probeServerStream(streamId); - - if (serverTarget) { - // pass the full server side identity (may carry a ::model suffix) so the GET routes - // straight to the owning session, no probe or fan out - await this.attachServerStream(convId, serverTarget.conversation_id); - - return; - } - - // fallback: local state remembers an interrupted byte offset for this conv, the server may - // still have a live session matching that identity (we just lost the bytes mid stream). retry - // with the frozen identity, the server probe inside attachServerStream tells us if it exists - if (!localState) { - return; - } - - // quiet status probe first: a full attach flips the loading UI on every try, probing - // keeps the retry loop invisible while the owning model is still loading (503) - const status = await ChatService.probeResumeStatus(streamId); - - if (status === 503) { - // make the wait visible: the empty assistant row persisted at send time renders - // the processing info, whose model load percentage flows from the models feed - this.resumePendingConvs.add(convId); - this.setChatLoading(convId, true); - - if (!this.resumeRetryTimers.has(convId)) { - this.resumeRetryTimers.set( - convId, - setTimeout(() => { - this.resumeRetryTimers.delete(convId); - void this.discoverActiveStream(convId); - }, STREAM_RESUME_RETRY_MS) - ); - } - - return; - } - - if (this.resumePendingConvs.delete(convId) && status !== 200) { - // the wait is over without a session to attach, drop the visible loading state - this.setChatLoading(convId, false); - } - - if (status === 0) { - // transient network failure, the next mount or visibility change retries - return; - } - - if (status !== 200) { - // the session is gone (stopped, TTL expired), nothing to resume anymore - ChatService.clearStreamState(convId); - - return; - } - - await this.attachServerStream(convId, streamId); - - // if attachServerStream failed (session gone, TTL expired), clear the local state to avoid retrying forever - if (!this.chatStreamingStates.has(convId) && !this.chatLoadingStates.get(convId)) { - ChatService.clearStreamState(convId); - } - } finally { - this.discoveringConvs.delete(convId); - } - } - - private findLastAssistantIdx(messages: DatabaseMessage[]): number { - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === MessageRole.ASSISTANT) return i; - } - - return -1; - } - - private findLastUserIdx(messages: DatabaseMessage[]): number { - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === MessageRole.USER) return i; - } - - return -1; - } - - clearUIState(): void { - this.isLoading = false; - this.currentResponse = ''; - this.isStreamingActive = false; - } - - setActiveProcessingConversation(conversationId: string | null): void { - this.activeConversationId = conversationId; - this.activeProcessingState = conversationId - ? this.processingStates.get(conversationId) || null - : null; - } - - getProcessingState(conversationId: string): ApiProcessingState | null { - return this.processingStates.get(conversationId) || null; - } - - private setProcessingState(conversationId: string, state: ApiProcessingState | null): void { - if (state === null) this.processingStates.delete(conversationId); - else this.processingStates.set(conversationId, state); - - if (conversationId === this.activeConversationId) this.activeProcessingState = state; - } - - clearProcessingState(conversationId: string): void { - this.processingStates.delete(conversationId); - - if (conversationId === this.activeConversationId) this.activeProcessingState = null; - } - - getActiveProcessingState(): ApiProcessingState | null { - return this.activeProcessingState; - } - - getCurrentProcessingStateSync(): ApiProcessingState | null { - return this.activeProcessingState; - } - - private setStreamingActive(active: boolean): void { - this.isStreamingActive = active; - } - - isStreaming(): boolean { - return this.isStreamingActive; - } - - private getOrCreateAbortController(convId: string): AbortController { - let c = this.abortControllers.get(convId); - - if (!c || c.signal.aborted) { - c = new AbortController(); - this.abortControllers.set(convId, c); - } - - return c; - } - - private abortRequest(convId?: string): void { - if (convId) { - const c = this.abortControllers.get(convId); - - if (c) { - c.abort(); - this.abortControllers.delete(convId); - } - } else { - for (const c of this.abortControllers.values()) c.abort(); - this.abortControllers.clear(); - } - } - - /** - * Abort the current agentic flow signal without clearing loading state. - * Used by "Send immediately" to force the agentic loop to exit so that - * the pending steering message can be re-sent. - * - * Any tool calls captured mid-stream are dropped before the abort so the - * pending message (or a manual follow-up) does not re-send a half-received - * tool call with invalid JSON arguments to the server. Mirrors what the - * Stop button already does through stopGenerationForChat. - */ - async abortCurrentFlow(convId: string): Promise { - await this.savePartialResponseIfNeeded(convId); - const c = this.abortControllers.get(convId); - - if (c) { - c.abort(); - this.abortControllers.delete(convId); - } - } - - private showErrorDialog(state: ErrorDialogState | null): void { - this.errorDialogState = state; - } - - dismissErrorDialog(): void { - this.errorDialogState = null; - } - - clearEditMode(): void { - this.isEditModeActive = false; - this.addFilesHandler = null; - } - - isEditing(): boolean { - return this.isEditModeActive; - } - - setEditModeActive(handler: (files: File[]) => void): void { - this.isEditModeActive = true; - this.addFilesHandler = handler; - } - - getAddFilesHandler(): ((files: File[]) => void) | null { - return this.addFilesHandler; - } - - clearPendingEditMessageId(): void { - this.pendingEditMessageId = null; - } - - savePendingDraft(message: string, files: ChatUploadedFile[]): void { - this._pendingDraftMessage = message; - this._pendingDraftFiles = [...files]; - } - - consumePendingDraft(): { message: string; files: ChatUploadedFile[] } | null { - if (!this._pendingDraftMessage && this._pendingDraftFiles.length === 0) return null; - - const d = { files: [...this._pendingDraftFiles], message: this._pendingDraftMessage }; - - this._pendingDraftMessage = ''; - this._pendingDraftFiles = []; - - return d; - } - - hasPendingDraft(): boolean { - return Boolean(this._pendingDraftMessage) || this._pendingDraftFiles.length > 0; - } - - getAllLoadingChats(): string[] { - // union of local (this browser is piping) and remote (backend reports a running session - // for this conv but no local pipe yet) sources. the sidebar shows one spinner per entry - const out = new SvelteSet(this.chatLoadingStates.keys()); - - for (const id of this.remoteRunningConvs) { - out.add(id); - } - - return Array.from(out); - } - - getAllStreamingChats(): string[] { - return Array.from(this.chatStreamingStates.keys()); - } - - /** - * Resync the remote running convs set from the backend. Called by the layout at mount and on - * visibilitychange, no polling. A snapshot semantic: the set is replaced wholesale, stale entries - * for sessions that finalized while the browser was elsewhere are dropped naturally. - */ - async syncRemoteRunningStreams(): Promise { - // the conversations store loads from IndexedDB asynchronously, the +layout onMount caller - // fires before that finishes. read ids straight from the DB so the result does not depend - // on the store init race, and the sidebar spinners light up at first paint for every conv - // the user owns even if it has not been hydrated into the store yet - let ids: string[]; - - try { - const all = await DatabaseService.getAllConversations(); - - ids = all.map((c) => c.id).filter((id) => !!id); - } catch (e) { - console.warn('syncRemoteRunningStreams DB read failed:', e); - - return; - } - - // only ask about conv ids the user already owns - if (ids.length === 0) { - for (const id of Array.from(this.remoteRunningConvs)) { - this.remoteRunningConvs.delete(id); - } - - return; - } - - // rebuild the frozen conv::model identity per conv so a session started with a model still - // matches. the server response is mapped back to the bare id below for the sidebar set - const lookupIds = ids.map((id) => - ChatService.resumeStreamIdentity(id, ChatService.getStreamState(id), null) - ); - - let sessions: ApiStreamSession[]; - - try { - sessions = await ChatService.lookupStreamSessions(lookupIds); - } catch (e) { - console.warn('syncRemoteRunningStreams lookup failed:', e); - - return; - } - const running = new SvelteSet(); - - for (const s of sessions) { - if (s && !s.is_done && typeof s.conversation_id === 'string' && s.conversation_id) { - // strip the optional ::model suffix, the sidebar set is keyed by the bare conv id - const sepIdx = s.conversation_id.indexOf(CONVERSATION_ID_SEPARATOR); - const bareId = sepIdx === -1 ? s.conversation_id : s.conversation_id.slice(0, sepIdx); - - running.add(bareId); - } - } - for (const id of Array.from(this.remoteRunningConvs)) { - if (!running.has(id)) { - this.remoteRunningConvs.delete(id); - } - } - for (const id of running) { - this.remoteRunningConvs.add(id); - } - } - - getChatStreaming(convId: string): { response: string; messageId: string } | undefined { - return this.getChatStreamingState(convId); - } - - isChatLoading(convId: string): boolean { - return this.chatLoadingStates.get(convId) || false; - } - - private isChatLoadingInternal(convId: string): boolean { - return this.chatLoadingStates.has(convId) || this.chatStreamingStates.has(convId); - } - - hasPendingMessage(convId: string): boolean { - return this._pendingMessages.has(convId); - } - - pendingMessageContent(convId: string): string | null { - return this._pendingMessages.get(convId)?.content ?? null; - } - - pendingMessageExtras(convId: string): DatabaseMessageExtra[] | undefined { - return this._pendingMessages.get(convId)?.extras; - } - - injectPendingMessage(convId: string, content: string, extras?: DatabaseMessageExtra[]): void { - this._pendingMessages.set(convId, { content, extras }); - } - - clearPendingMessage(convId: string): void { - this._pendingMessages.delete(convId); - } - - consumePendingMessage( - convId: string - ): { content: string; extras?: DatabaseMessageExtra[] } | null { - const msg = this._pendingMessages.get(convId); - - if (!msg) return null; - - this._pendingMessages.delete(convId); - - return msg; - } - - private touchConversationState(convId: string): void { - this.conversationStateTimestamps.set(convId, { lastAccessed: Date.now() }); - } - - cleanupOldConversationStates(activeConversationIds?: string[]): number { - const now = Date.now(); - const activeIdsList = activeConversationIds ?? []; - const preserveIds = this.activeConversationId - ? [...activeIdsList, this.activeConversationId] - : activeIdsList; - const allConvIds = [ - ...new Set([ - ...this.chatLoadingStates.keys(), - ...this.chatStreamingStates.keys(), - ...this.abortControllers.keys(), - ...this.processingStates.keys(), - ...this.conversationStateTimestamps.keys() - ]) - ]; - const cleanupCandidates: Array<{ convId: string; lastAccessed: number }> = []; - - for (const convId of allConvIds) { - if (preserveIds.includes(convId)) continue; - - if (this.chatLoadingStates.get(convId)) continue; - - if (this.chatStreamingStates.has(convId)) continue; - - const ts = this.conversationStateTimestamps.get(convId); - - cleanupCandidates.push({ convId, lastAccessed: ts?.lastAccessed ?? 0 }); - } - cleanupCandidates.sort((a, b) => a.lastAccessed - b.lastAccessed); - let cleanedUp = 0; - - for (const { convId, lastAccessed } of cleanupCandidates) { - if ( - cleanupCandidates.length - cleanedUp > INACTIVE_CONVERSATION.MAX_STATES || - now - lastAccessed > INACTIVE_CONVERSATION.MAX_AGE_MS - ) { - this.cleanupConversationState(convId); - cleanedUp++; - } - } - - return cleanedUp; - } - private cleanupConversationState(convId: string): void { - const c = this.abortControllers.get(convId); - - if (c && !c.signal.aborted) c.abort(); - - this.chatLoadingStates.delete(convId); - this.chatStreamingStates.delete(convId); - this.abortControllers.delete(convId); - this.processingStates.delete(convId); - this.conversationStateTimestamps.delete(convId); - } - getTrackedConversationCount(): number { - return new Set([ - ...this.chatLoadingStates.keys(), - ...this.chatStreamingStates.keys(), - ...this.abortControllers.keys(), - ...this.processingStates.keys() - ]).size; - } - - private getMessageByIdWithRole( - messageId: string, - expectedRole?: MessageRole - ): { message: DatabaseMessage; index: number } | null { - const index = conversationsStore.findMessageIndex(messageId); - - if (index === -1) return null; - - const message = conversationsStore.activeMessages[index]; - - if (expectedRole && message.role !== expectedRole) return null; - - return { index, message }; - } - - async addMessage( - role: MessageRole, - content: string, - type: MessageType = MessageType.TEXT, - parent: string = '-1', - extras?: DatabaseMessageExtra[], - isSynthetic?: boolean - ): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) throw new Error('No active conversation'); - - let parentId: string | null = null; - - if (parent === '-1') { - const am = conversationsStore.activeMessages; - - if (am.length > 0) parentId = am[am.length - 1].id; - else { - const all = await conversationsStore.getConversationMessages(activeConv.id); - const r = all.find((m) => m.parent === null && m.type === 'root'); - - parentId = r ? r.id : await DatabaseService.createRootMessage(activeConv.id); - } - } else parentId = parent; - - const message = await DatabaseService.createMessageBranch( - { - children: [], - content, - convId: activeConv.id, - extra: extras, - isSynthetic, - role, - timestamp: Date.now(), - toolCalls: '', - type - }, - parentId - ); - - conversationsStore.addMessageToActive(message); - await conversationsStore.updateCurrentNode(message.id); - conversationsStore.updateConversationTimestamp(); - - return message; - } - - /** - * Record a working-directory change into chat history as a synthetic - * user message, so the model sees it on its next turn (the client - * sends the cwd itself via the x-tool-cwd header on tool calls). - * A plain user message is used because some chat templates reject - * tool messages without a preceding tool call. - */ - async recordCwdChange(cwd: string | null): Promise { - const content = cwd - ? formatCwdMessage(cwd, await toolsStore.resolveServerHome()) - : CWD_CLEARED_TEXT; - // Reuse the trailing cwd row when it is already the last message, so - // repeated picks update it in place instead of stacking another row. - const last = conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1]; - - if (last && last.role === MessageRole.USER && last.isSynthetic === true) { - await DatabaseService.updateMessage(last.id, { content, isSynthetic: true }); - conversationsStore.updateMessageAtIndex(conversationsStore.activeMessages.length - 1, { - content, - isSynthetic: true - }); - - return; - } - - await this.addMessage(MessageRole.USER, content, MessageType.TEXT, '-1', undefined, true); - } - - async addSystemPrompt(): Promise { - let activeConv = conversationsStore.activeConversation; - - if (!activeConv) { - await conversationsStore.createConversation(); - activeConv = conversationsStore.activeConversation; - } - - if (!activeConv) return; - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); - const rootId = rootMessage - ? rootMessage.id - : await DatabaseService.createRootMessage(activeConv.id); - const existingSystemMessage = allMessages.find( - (m) => m.role === MessageRole.SYSTEM && m.parent === rootId - ); - - if (existingSystemMessage) { - this.pendingEditMessageId = existingSystemMessage.id; - - if (!conversationsStore.activeMessages.some((m) => m.id === existingSystemMessage.id)) - conversationsStore.activeMessages.unshift(existingSystemMessage); - - return; - } - - const am = conversationsStore.activeMessages; - const firstActiveMessage = am.find((m) => m.parent === rootId); - const systemMessage = await DatabaseService.createSystemMessage( - activeConv.id, - SYSTEM_MESSAGE_PLACEHOLDER, - rootId - ); - - if (firstActiveMessage) { - await DatabaseService.updateMessage(firstActiveMessage.id, { - parent: systemMessage.id - }); - await DatabaseService.updateMessage(systemMessage.id, { - children: [firstActiveMessage.id] - }); - const updatedRootChildren = rootMessage - ? rootMessage.children.filter((id: string) => id !== firstActiveMessage.id) - : []; - - await DatabaseService.updateMessage(rootId, { - children: [ - ...updatedRootChildren.filter((id: string) => id !== systemMessage.id), - systemMessage.id - ] - }); - const firstMsgIndex = conversationsStore.findMessageIndex(firstActiveMessage.id); - - if (firstMsgIndex !== -1) - conversationsStore.updateMessageAtIndex(firstMsgIndex, { - parent: systemMessage.id - }); - } - - conversationsStore.activeMessages.unshift(systemMessage); - this.pendingEditMessageId = systemMessage.id; - conversationsStore.updateConversationTimestamp(); - } catch (error) { - console.error('Failed to add system prompt:', error); - } - } - - async removeSystemPromptPlaceholder(messageId: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return false; - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const systemMessage = findMessageById(allMessages, messageId); - - if (!systemMessage || systemMessage.role !== MessageRole.SYSTEM) return false; - - const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); - - if (!rootMessage) return false; - - if (allMessages.length === 2 && systemMessage.children.length === 0) { - await conversationsStore.deleteConversation(activeConv.id); - - return true; - } - - for (const childId of systemMessage.children) { - await DatabaseService.updateMessage(childId, { parent: rootMessage.id }); - const childIndex = conversationsStore.findMessageIndex(childId); - - if (childIndex !== -1) - conversationsStore.updateMessageAtIndex(childIndex, { parent: rootMessage.id }); - } - await DatabaseService.updateMessage(rootMessage.id, { - children: [ - ...rootMessage.children.filter((id: string) => id !== messageId), - ...systemMessage.children - ] - }); - await DatabaseService.deleteMessage(messageId); - const systemIndex = conversationsStore.findMessageIndex(messageId); - - if (systemIndex !== -1) conversationsStore.activeMessages.splice(systemIndex, 1); - - conversationsStore.updateConversationTimestamp(); - - return false; - } catch (error) { - console.error('Failed to remove system prompt placeholder:', error); - - return false; - } - } - - private async createAssistantMessage(parentId?: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) throw new Error('No active conversation'); - - return await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId: activeConv.id, - model: null, - role: MessageRole.ASSISTANT, - timestamp: Date.now(), - toolCalls: '', - type: MessageType.TEXT - }, - parentId || null - ); - } - - async sendMessage(content: string, extras?: DatabaseMessageExtra[]): Promise { - if (!content.trim() && (!extras || extras.length === 0)) return; - - const activeConv = conversationsStore.activeConversation; - - // If agentic loop is running, inject as a steering message instead of starting a new flow - if (activeConv && agenticStore.isRunning(activeConv.id)) { - agenticStore.injectSteeringMessage(activeConv.id, content, extras); - - return; - } - - // If non-agentic streaming is active, queue as a pending message to send after completion - if (activeConv && this.isChatLoadingInternal(activeConv.id)) { - this.injectPendingMessage(activeConv.id, content, extras); - - return; - } - - // Cancel any in-flight pre-encode request - this.cancelPreEncode(); - - // Consume MCP resource attachments - converts them to extras and clears the live store - const resourceExtras = mcpStore.consumeResourceAttachmentsAsExtras(); - const allExtras = resourceExtras.length > 0 ? [...(extras || []), ...resourceExtras] : extras; - - let isNewConversation = false; - - if (!activeConv) { - await conversationsStore.createConversation(); - isNewConversation = true; - } - - const currentConv = conversationsStore.activeConversation; - - if (!currentConv) return; - - this.showErrorDialog(null); - this.setChatLoading(currentConv.id, true); - this.clearChatStreaming(currentConv.id); - try { - let parentIdForUserMessage: string | undefined; - - if (isNewConversation) { - const rootId = await DatabaseService.createRootMessage(currentConv.id); - const currentConfig = settingsStore.config; - const systemPrompt = currentConfig.systemMessage?.toString().trim(); - - let sysOrRootId = rootId; - - if (systemPrompt) { - const systemMessage = await DatabaseService.createSystemMessage( - currentConv.id, - systemPrompt, - rootId - ); - - conversationsStore.addMessageToActive(systemMessage); - sysOrRootId = systemMessage.id; - } - - // Reflect a working directory picked on the new-chat screen into - // chat history before the first user message, so the model sees - // it on its first turn. createConversation() has already threaded - // the pending pick onto the conversation. - if (currentConv.cwd) { - const cwdMessage = await this.addMessage( - MessageRole.USER, - formatCwdMessage(currentConv.cwd, await toolsStore.resolveServerHome()), - MessageType.TEXT, - sysOrRootId, - undefined, - true - ); - - parentIdForUserMessage = cwdMessage.id; - } else { - parentIdForUserMessage = sysOrRootId; - } - } - - const userMessage = await this.addMessage( - MessageRole.USER, - content, - MessageType.TEXT, - parentIdForUserMessage ?? '-1', - allExtras - ); - - if (isNewConversation && content) - await conversationsStore.updateConversationName( - currentConv.id, - generateConversationTitle( - content, - Boolean(settingsStore.config.titleGenerationUseFirstLine) - ) - ); - - const assistantMessage = await this.createAssistantMessage(userMessage.id); - - conversationsStore.addMessageToActive(assistantMessage); - await this.streamChatCompletion( - conversationsStore.activeMessages.slice(0, -1), - assistantMessage, - undefined, - undefined, - undefined, - settingsStore.config.titleGenerationUseLLM && isNewConversation ? content : undefined - ); - } catch (error) { - if (isAbortError(error)) { - this.setChatLoading(currentConv.id, false); - - return; - } - - console.error('Failed to send message:', error); - this.setChatLoading(currentConv.id, false); - const dialogType = - error instanceof Error && error.name === 'TimeoutError' - ? ErrorDialogType.TIMEOUT - : ErrorDialogType.SERVER; - const contextInfo = ( - error as Error & { contextInfo?: { n_prompt_tokens: number; n_ctx: number } } - ).contextInfo; - - this.showErrorDialog({ - contextInfo, - message: error instanceof Error ? error.message : 'Unknown error', - type: dialogType - }); - } - } - - private async streamChatCompletion( - allMessages: DatabaseMessage[], - assistantMessage: DatabaseMessage, - onComplete?: (content: string) => Promise, - onError?: (error: Error) => void, - modelOverride?: string | null, - firstUserMessageContent?: string - ): Promise { - // the ::model suffix in the stream identity is only for router mode, where it routes to the - // owning child. in single-model mode the identity stays the bare conv id so that attach, stop - // and reattach all agree, regardless of fresh send vs regenerate passing a resolved model - let effectiveModel: string | null | undefined = undefined; - - if (serverStore.isRouterMode) { - const conversationModel = getConversationModel(allMessages); - - effectiveModel = modelOverride || modelsStore.selectedModelName || conversationModel; - } - - if (serverStore.isRouterMode && effectiveModel) { - if (!modelsStore.getModelProps(effectiveModel)) - await modelsStore.fetchModelProps(effectiveModel); - } - - // Mutable state for the current message being streamed - let currentMessageId = assistantMessage.id; - let streamedContent = ''; - let streamedReasoningContent = ''; - let resolvedModel: string | null = null; - let modelPersisted = false; - - const convId = assistantMessage.convId; - - // Tracks the last message created in this flow. Used as the parent for the next - // turn's assistant message so createAssistantMessage does not have to read - // conversationsStore.activeMessages, which may belong to a different conversation - // after the user navigates while the loop is still running. - let lastCreatedInFlow = currentMessageId; - - // freeze the POST identity from t0 so a stop cancels with the exact session key, - // never a stale or empty model resolved later - this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); - - const recordModel = (modelName: string | null | undefined, persistImmediately = true): void => { - if (!modelName) return; - - const n = normalizeModelName(modelName); - - if (!n || n === resolvedModel) return; - - resolvedModel = n; - const idx = conversationsStore.findMessageIndex(currentMessageId); - - conversationsStore.updateMessageAtIndex(idx, { model: n }); - - if (persistImmediately && !modelPersisted) { - modelPersisted = true; - DatabaseService.updateMessage(currentMessageId, { model: n }).catch(() => { - modelPersisted = false; - resolvedModel = null; - }); - } - }; - - let completionIdRecorded = false; - - const recordCompletionId = (id: string): void => { - if (!id || completionIdRecorded) return; - - completionIdRecorded = true; - const idx = conversationsStore.findMessageIndex(currentMessageId); - - conversationsStore.updateMessageAtIndex(idx, { completionId: id }); - DatabaseService.updateMessage(currentMessageId, { completionId: id }).catch(() => { - completionIdRecorded = false; - }); - }; - const updateStreamingUI = () => { - this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); - const idx = conversationsStore.findMessageIndex(currentMessageId); - - conversationsStore.updateMessageAtIndex(idx, { content: streamedContent }); - }; - const cleanupStreamingState = () => { - this.setStreamingActive(false); - this.setChatLoading(convId, false); - this.clearChatStreaming(convId, currentMessageId); - this.setProcessingState(convId, null); - }; - - this.setStreamingActive(true); - this.setActiveProcessingConversation(convId); - const abortController = this.getOrCreateAbortController(convId); - const streamCallbacks: ChatStreamCallbacks = { - createAssistantMessage: async () => { - // Reset streaming state for new message - streamedContent = ''; - streamedReasoningContent = ''; - - const msg = await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId, - model: resolvedModel, - role: MessageRole.ASSISTANT, - timestamp: Date.now(), - toolCalls: '', - type: MessageType.TEXT - }, - lastCreatedInFlow - ); - - if (conversationsStore.activeConversation?.id === convId) { - conversationsStore.addMessageToActive(msg); - } - - currentMessageId = msg.id; - lastCreatedInFlow = msg.id; - - return msg; - }, - createToolResultMessage: async ( - toolCallId: string, - content: string, - extras?: DatabaseMessageExtra[], - toolCwd?: string - ) => { - const msg = await DatabaseService.createMessageBranch( - { - children: [], - content, - convId, - extra: extras, - role: MessageRole.TOOL, - timestamp: Date.now(), - toolCallId, - toolCalls: '', - toolCwd, - type: MessageType.TEXT - }, - currentMessageId - ); - - // mirror into the active store and move the node pointer only when this - // conversation is displayed; otherwise persist the node move straight to - // the db for the owning conv so a foreign conv's currNode stays untouched - if (conversationsStore.activeConversation?.id === convId) { - conversationsStore.addMessageToActive(msg); - await conversationsStore.updateCurrentNode(msg.id); - } else { - await DatabaseService.updateCurrentNode(convId, msg.id); - } - - lastCreatedInFlow = msg.id; - - return msg; - }, - onAssistantTurnComplete: async ( - content: string, - reasoningContent: string | undefined, - timings: ChatMessageTimings | undefined, - toolCalls: import('$lib/types/api').ApiChatCompletionToolCall[] | undefined - ) => { - const updateData: Record = { - content, - reasoningContent: reasoningContent || undefined, - timings, - toolCalls: toolCalls ? JSON.stringify(toolCalls) : '' - }; - - if (resolvedModel && !modelPersisted) updateData.model = resolvedModel; - - await DatabaseService.updateMessage(currentMessageId, updateData); - const idx = conversationsStore.findMessageIndex(currentMessageId); - const uiUpdate: Partial = { - content, - reasoningContent: reasoningContent || undefined, - toolCalls: toolCalls ? JSON.stringify(toolCalls) : '' - }; - - if (timings) uiUpdate.timings = timings; - - if (resolvedModel) uiUpdate.model = resolvedModel; - - // touch the active ui array and node pointer only when this conversation - // is displayed; otherwise persist the node move straight to the db so a - // foreign conv's currNode stays untouched - if (conversationsStore.activeConversation?.id === convId) { - conversationsStore.updateMessageAtIndex(idx, uiUpdate); - await conversationsStore.updateCurrentNode(currentMessageId); - } else { - await DatabaseService.updateCurrentNode(convId, currentMessageId); - } - }, - onAttachments: (messageId: string, extras: DatabaseMessageExtra[]) => { - if (!extras.length) return; - - const idx = conversationsStore.findMessageIndex(messageId); - - if (idx === -1) return; - - const msg = conversationsStore.activeMessages[idx]; - const updatedExtras = [...(msg.extra || []), ...extras]; - - conversationsStore.updateMessageAtIndex(idx, { extra: updatedExtras }); - DatabaseService.updateMessage(messageId, { extra: updatedExtras }).catch(console.error); - }, - onChunk: (chunk: string) => { - streamedContent += chunk; - updateStreamingUI(); - this.setChatReasoning(convId, false); - }, - onCompletionId: (id: string) => recordCompletionId(id), - onError: async (error: Error) => { - this.setStreamingActive(false); - - if (isAbortError(error)) { - cleanupStreamingState(); - // If aborted with a pending message (e.g. "Send immediately"), re-send it - const pending = this.consumePendingMessage(convId); - - if (pending) { - this.sendMessage(pending.content, pending.extras); - } - - return; - } - - console.error('Streaming error:', error); - // keep whatever was streamed so far, the message stays in memory and in DB - await this.savePartialResponseIfNeeded(convId); - cleanupStreamingState(); - this.clearPendingMessage(convId); - - const contextInfo = ( - error as Error & { contextInfo?: { n_prompt_tokens: number; n_ctx: number } } - ).contextInfo; - - this.showErrorDialog({ - contextInfo, - message: error.message, - type: error.name === 'TimeoutError' ? ErrorDialogType.TIMEOUT : ErrorDialogType.SERVER - }); - - if (onError) onError(error); - }, - onFlowComplete: (finalTimings?: ChatMessageTimings) => { - if (finalTimings) { - const idx = conversationsStore.findMessageIndex(assistantMessage.id); - - conversationsStore.updateMessageAtIndex(idx, { timings: finalTimings }); - DatabaseService.updateMessage(assistantMessage.id, { - timings: finalTimings - }).catch(console.error); - } - - cleanupStreamingState(); - - if (onComplete) onComplete(streamedContent); - - if (serverStore.isRouterMode) modelsStore.fetchRouterModels().catch(console.error); - - // Pre-encode conversation in KV cache for faster next turn - if (settingsStore.config.preEncodeConversation) { - this.triggerPreEncode( - allMessages, - assistantMessage, - streamedContent, - effectiveModel, - !!settingsStore.config.excludeReasoningFromContext - ); - } - }, - onModel: (modelName: string) => recordModel(modelName), - onReasoningChunk: (chunk: string) => { - streamedReasoningContent += chunk; - // mark streaming state so a stop mid-thinking can persist the partial reasoning - this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); - const idx = conversationsStore.findMessageIndex(currentMessageId); - - conversationsStore.updateMessageAtIndex(idx, { - reasoningContent: streamedReasoningContent - }); - this.setChatReasoning(convId, true); - }, - onTimings: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => { - const tokensPerSecond = - timings?.predicted_ms && timings?.predicted_n - ? (timings.predicted_n / timings.predicted_ms) * 1000 - : 0; - - this.updateProcessingStateFromTimings( - { - cache_n: timings?.cache_n || 0, - predicted_n: timings?.predicted_n || 0, - predicted_per_second: tokensPerSecond, - prompt_ms: timings?.prompt_ms, - prompt_n: timings?.prompt_n || 0, - prompt_progress: promptProgress - }, - convId - ); - }, - onToolCallsStreaming: (toolCalls) => { - const idx = conversationsStore.findMessageIndex(currentMessageId); - - conversationsStore.updateMessageAtIndex(idx, { - toolCalls: JSON.stringify(toolCalls) - }); - }, - onTurnComplete: (intermediateTimings: ChatMessageTimings) => { - // Update the first assistant message with cumulative agentic timings - const idx = conversationsStore.findMessageIndex(assistantMessage.id); - - conversationsStore.updateMessageAtIndex(idx, { timings: intermediateTimings }); - }, - updateToolResultMessage: async ( - messageId: string, - content: string, - extras?: DatabaseMessageExtra[] - ) => { - // Persist latest content + merged extras; mirror into the active - // store so the chat view sees live updates for streaming tools - // (e.g. exec_shell_command). The existing tool message node - // pointer stays put - the renderer is already scoped to it. - const updates: Partial = { content }; - - if (extras) { - const idx = conversationsStore.findMessageIndex(messageId); - const existing = idx >= 0 ? (conversationsStore.activeMessages[idx]?.extra ?? []) : []; - const merged = [...existing, ...extras]; - - updates.extra = merged; - } - - if (conversationsStore.activeConversation?.id === convId) { - const idx = conversationsStore.findMessageIndex(messageId); - - if (idx >= 0) conversationsStore.updateMessageAtIndex(idx, updates); - } - - await DatabaseService.updateMessage(messageId, updates); - } - }; - const perChatOverrides = conversationsStore.getAllMcpServerOverrides(); - - { - const agenticResult = await agenticStore.runAgenticFlow({ - callbacks: streamCallbacks, - conversationId: convId, - flowRootMessageId: assistantMessage.id, - messages: allMessages, - options: { - ...this.getApiOptions(), - ...(effectiveModel ? { model: effectiveModel } : {}) - }, - perChatOverrides, - signal: abortController.signal - }); - - if (agenticResult.handled) { - // Generate LLM based title for new conversations after agentic flow completes - if (firstUserMessageContent) { - await this.generateTitleWithLLM(firstUserMessageContent, streamedContent, convId); - } - - // Check if there's a pending steering message to re-send - const pending = agenticStore.consumePendingSteeringMessage(convId); - - if (pending) { - await this.sendMessage(pending.content, pending.extras); - } - - return; - } - } - - await ChatService.sendMessage( - allMessages, - { - ...this.getApiOptions(), - ...(effectiveModel ? { model: effectiveModel } : {}), - onChunk: streamCallbacks.onChunk, - onComplete: async ( - finalContent?: string, - reasoningContent?: string, - timings?: ChatMessageTimings, - toolCalls?: string - ) => { - const content = streamedContent || finalContent || ''; - const reasoning = streamedReasoningContent || reasoningContent; - const updateData: Record = { - content, - reasoningContent: reasoning || undefined, - timings, - toolCalls: toolCalls || '' - }; - - if (resolvedModel && !modelPersisted) updateData.model = resolvedModel; - - await DatabaseService.updateMessage(currentMessageId, updateData); - const idx = conversationsStore.findMessageIndex(currentMessageId); - const uiUpdate: Partial = { - content, - reasoningContent: reasoning || undefined, - toolCalls: toolCalls || '' - }; - - if (timings) uiUpdate.timings = timings; - - if (resolvedModel) uiUpdate.model = resolvedModel; - - conversationsStore.updateMessageAtIndex(idx, uiUpdate); - await conversationsStore.updateCurrentNode(currentMessageId); - cleanupStreamingState(); - - if (onComplete) await onComplete(content); - - if (serverStore.isRouterMode) modelsStore.fetchRouterModels().catch(console.error); - - // Generate LLM based title for new conversations (avoids stale reference - // issue when user switches conversations while streaming) - if (firstUserMessageContent) { - await this.generateTitleWithLLM(firstUserMessageContent, streamedContent, convId); - } - - // Check if there's a pending message queued during streaming - const pending = this.consumePendingMessage(convId); - - if (pending) { - await this.sendMessage(pending.content, pending.extras); - } - }, - onCompletionId: streamCallbacks.onCompletionId, - onConnectionState: (state: StreamConnectionState) => { - if (convId === conversationsStore.activeConversation?.id) { - this.streamConnectionState = state; - } - }, - onError: streamCallbacks.onError, - onModel: streamCallbacks.onModel, - onReasoningChunk: streamCallbacks.onReasoningChunk, - onTimings: streamCallbacks.onTimings, - stream: true - }, - convId, - abortController.signal - ); - } - - async stopGeneration(): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - await this.stopGenerationForChat(activeConv.id); - } - async stopGenerationForChat(convId: string): Promise { - await this.savePartialResponseIfNeeded(convId); - this.setStreamingActive(false); - // tell the server to stop the generation, not just drop the HTTP socket. without this the - // detached drain keeps producing tokens until eos or max_tokens. use the frozen identity - // captured when the session started, not the live dropdown - const streamStateForStop = this.chatStreamingStates.get(convId); - const modelForStop = streamStateForStop?.model ?? ChatService.getStreamState(convId)?.model; - - void ChatService.cancelServerStream(convId, modelForStop); - // an explicit stop leaves nothing to resume and kills a pending resume retry - ChatService.clearStreamState(convId); - const retryTimer = this.resumeRetryTimers.get(convId); - - if (retryTimer !== undefined) { - clearTimeout(retryTimer); - this.resumeRetryTimers.delete(convId); - } - - this.resumePendingConvs.delete(convId); - this.abortRequest(convId); - this.setChatLoading(convId, false); - this.clearChatStreaming(convId); - this.setProcessingState(convId, null); - this.clearPendingMessage(convId); - } - - private async generateTitleWithLLM( - userContent: string, - assistantContent: string, - convId: string - ): Promise { - const effectiveModel = - serverStore.isRouterMode && modelsStore.selectedModelName - ? modelsStore.selectedModelName - : undefined; - const configValue = settingsStore.config; - const titlePromptTemplate = - typeof configValue.titleGenerationPrompt === 'string' && - configValue.titleGenerationPrompt.trim() - ? configValue.titleGenerationPrompt - : TITLE_GENERATION.DEFAULT_PROMPT; - const titlePrompt = titlePromptTemplate - .replace('{{USER}}', String(userContent || '')) - .replace('{{ASSISTANT}}', String(assistantContent || '')); - const titleMessage: ApiChatMessageData = { - content: titlePrompt, - role: MessageRole.USER - }; - const titleResponse = await ChatService.generateTitle(titleMessage, effectiveModel); - - if (!titleResponse) { - return; - } - - let cleanTitle = titleResponse.trim(); - - cleanTitle = cleanTitle - .replace(TITLE_GENERATION.PREFIX_PATTERN, '') - .replace(TITLE_GENERATION.QUOTE_PATTERN, '') - .trim(); - - if (!cleanTitle || cleanTitle.length < TITLE_GENERATION.MIN_LENGTH) { - const firstLine = userContent.split('\n').find((l) => l.trim().length > 0); - - cleanTitle = firstLine ? firstLine.trim() : TITLE_GENERATION.FALLBACK; - } - - if (cleanTitle && cleanTitle.length >= TITLE_GENERATION.MIN_LENGTH) { - await conversationsStore.updateConversationName(convId, cleanTitle); - } - } - - private async savePartialResponseIfNeeded(convId?: string): Promise { - const conversationId = convId || conversationsStore.activeConversation?.id; - - if (!conversationId) return; - - const streamingState = this.getChatStreamingState(conversationId); - - if (!streamingState) return; - - const messages = - conversationId === conversationsStore.activeConversation?.id - ? conversationsStore.activeMessages - : await conversationsStore.getConversationMessages(conversationId); - - if (!messages.length) return; - - const lastMessage = messages[messages.length - 1]; - - if (lastMessage?.role !== MessageRole.ASSISTANT) return; - - const partialContent = streamingState.response; - const partialReasoning = lastMessage.reasoningContent || ''; - // snapshot the streamed tool calls before clearing so we still know whether - // anything was captured when deciding to skip the DB write below - const hadPartialToolCalls = !!lastMessage.toolCalls?.trim(); - - // nothing to persist when content, reasoning, and streamed tool calls are all empty - // (e.g. stop before any token). otherwise drop the partial tool call and write whatever - // was streamed: incomplete arguments (truncated JSON, missing closing quote) would - // otherwise be re-sent to the server on the next turn and rejected. - if (!partialContent.trim() && !partialReasoning.trim() && !hadPartialToolCalls) return; - - try { - const updateData: { - content?: string; - reasoningContent?: string; - toolCalls?: string; - timings?: ChatMessageTimings; - } = { - toolCalls: '' - }; - - if (partialContent.trim()) updateData.content = partialContent; - - if (partialReasoning.trim()) updateData.reasoningContent = partialReasoning; - - const lastKnownState = this.getProcessingState(conversationId); - - if (lastKnownState) { - updateData.timings = { - cache_n: lastKnownState.cacheTokens || 0, - predicted_ms: - lastKnownState.tokensPerSecond && lastKnownState.tokensDecoded - ? (lastKnownState.tokensDecoded / lastKnownState.tokensPerSecond) * 1000 - : undefined, - predicted_n: lastKnownState.tokensDecoded || 0, - prompt_ms: lastKnownState.promptMs, - prompt_n: lastKnownState.promptTokens || 0 - }; - } - - await DatabaseService.updateMessage(lastMessage.id, updateData); - lastMessage.content = partialContent; - // mirror the drop into the in-memory message so the next request sent via - // sendMessage (queued pending, Send immediately, or manual follow-up) reads - // the cleared value, not whatever the streaming widget had been showing - lastMessage.toolCalls = ''; - - if (updateData.timings) lastMessage.timings = updateData.timings; - } catch (error) { - lastMessage.content = partialContent; - lastMessage.toolCalls = ''; - console.error('Failed to save partial response:', error); - } - } - - async updateMessage(messageId: string, newContent: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - if (this.isChatLoadingInternal(activeConv.id)) await this.stopGeneration(); - - const result = this.getMessageByIdWithRole(messageId, MessageRole.USER); - - if (!result) return; - - const { index: messageIndex, message: messageToUpdate } = result; - const originalContent = messageToUpdate.content; - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); - const isFirstUserMessage = rootMessage && messageToUpdate.parent === rootMessage.id; - - conversationsStore.updateMessageAtIndex(messageIndex, { content: newContent }); - await DatabaseService.updateMessage(messageId, { content: newContent }); - - if (isFirstUserMessage && newContent.trim()) - await conversationsStore.updateConversationName( - activeConv.id, - generateConversationTitle( - newContent, - Boolean(settingsStore.config.titleGenerationUseFirstLine) - ) - ); - - const messagesToRemove = conversationsStore.activeMessages.slice(messageIndex + 1); - - if (messagesToRemove.length > 0) - await DatabaseService.deleteMessageCascading(activeConv.id, messagesToRemove[0].id); - - conversationsStore.sliceActiveMessages(messageIndex + 1); - conversationsStore.updateConversationTimestamp(); - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - const assistantMessage = await this.createAssistantMessage(); - - conversationsStore.addMessageToActive(assistantMessage); - await conversationsStore.updateCurrentNode(assistantMessage.id); - await this.streamChatCompletion( - conversationsStore.activeMessages.slice(0, -1), - assistantMessage, - undefined, - () => { - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(messageId), { - content: originalContent - }); - } - ); - } catch (error) { - if (!isAbortError(error)) console.error('Failed to update message:', error); - } - } - - async regenerateMessage(messageId: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; - - this.cancelPreEncode(); - const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); - - if (!result) return; - - const { index: messageIndex } = result; - - try { - const messagesToRemove = conversationsStore.activeMessages.slice(messageIndex); - - await DatabaseService.deleteMessageCascading(activeConv.id, messagesToRemove[0].id); - conversationsStore.sliceActiveMessages(messageIndex); - conversationsStore.updateConversationTimestamp(); - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - const parentMessageId = - conversationsStore.activeMessages.length > 0 - ? conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1].id - : undefined; - const assistantMessage = await this.createAssistantMessage(parentMessageId); - - conversationsStore.addMessageToActive(assistantMessage); - await this.streamChatCompletion( - conversationsStore.activeMessages.slice(0, -1), - assistantMessage - ); - } catch (error) { - if (!isAbortError(error)) console.error('Failed to regenerate message:', error); - - this.setChatLoading(activeConv?.id || '', false); - } - } - - async regenerateMessageWithBranching(messageId: string, modelOverride?: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; - - this.cancelPreEncode(); - try { - const idx = conversationsStore.findMessageIndex(messageId); - - if (idx === -1) return; - - const msg = conversationsStore.activeMessages[idx]; - - if (msg.role !== MessageRole.ASSISTANT) return; - - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const parentMessage = findMessageById(allMessages, msg.parent); - - if (!parentMessage) return; - - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - const newAssistantMessage = await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId: msg.convId, - model: null, - role: msg.role, - timestamp: Date.now(), - toolCalls: '', - type: msg.type - }, - parentMessage.id - ); - - await conversationsStore.updateCurrentNode(newAssistantMessage.id); - conversationsStore.updateConversationTimestamp(); - await conversationsStore.refreshActiveMessages(); - const conversationPath = filterByLeafNodeId( - allMessages, - parentMessage.id, - false - ) as DatabaseMessage[]; - const modelToUse = modelOverride || msg.model || undefined; - - await this.streamChatCompletion( - conversationPath, - newAssistantMessage, - undefined, - undefined, - modelToUse - ); - } catch (error) { - if (!isAbortError(error)) - console.error('Failed to regenerate message with branching:', error); - - this.setChatLoading(activeConv?.id || '', false); - } - } - - async getDeletionInfo(messageId: string): Promise<{ - totalCount: number; - userMessages: number; - assistantMessages: number; - messageTypes: string[]; - }> { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) - return { assistantMessages: 0, messageTypes: [], totalCount: 0, userMessages: 0 }; - - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const messageToDelete = findMessageById(allMessages, messageId); - - // For system messages, don't count descendants as they will be preserved (reparented to root) - if (messageToDelete?.role === MessageRole.SYSTEM) { - const messagesToDelete = allMessages.filter((m) => m.id === messageId); - - let assistantMessages = 0, - userMessages = 0; - - const messageTypes: string[] = []; - - for (const msg of messagesToDelete) { - if (msg.role === MessageRole.USER) { - userMessages++; - - if (!messageTypes.includes('user message')) messageTypes.push('user message'); - } else if (msg.role === MessageRole.ASSISTANT) { - assistantMessages++; - - if (!messageTypes.includes('assistant response')) messageTypes.push('assistant response'); - } - } - - return { assistantMessages, messageTypes, totalCount: 1, userMessages }; - } - - const descendants = findDescendantMessages(allMessages, messageId); - const allToDelete = [messageId, ...descendants]; - const messagesToDelete = allMessages.filter((m) => allToDelete.includes(m.id)); - - let assistantMessages = 0, - userMessages = 0; - - const messageTypes: string[] = []; - - for (const msg of messagesToDelete) { - if (msg.role === MessageRole.USER) { - userMessages++; - - if (!messageTypes.includes('user message')) messageTypes.push('user message'); - } else if (msg.role === MessageRole.ASSISTANT) { - assistantMessages++; - - if (!messageTypes.includes('assistant response')) messageTypes.push('assistant response'); - } - } - - return { assistantMessages, messageTypes, totalCount: allToDelete.length, userMessages }; - } - - async deleteMessage(messageId: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const messageToDelete = findMessageById(allMessages, messageId); - - if (!messageToDelete) return; - - const currentPath = filterByLeafNodeId(allMessages, activeConv.currNode || '', false); - const isInCurrentPath = currentPath.some((m) => m.id === messageId); - - if (isInCurrentPath && messageToDelete.parent) { - const siblings = allMessages.filter( - (m) => m.parent === messageToDelete.parent && m.id !== messageId - ); - - if (siblings.length > 0) { - const latestSibling = siblings.reduce((latest, sibling) => - sibling.timestamp > latest.timestamp ? sibling : latest - ); - - await conversationsStore.updateCurrentNode(findLeafNode(allMessages, latestSibling.id)); - } else if (messageToDelete.parent) { - await conversationsStore.updateCurrentNode( - findLeafNode(allMessages, messageToDelete.parent) - ); - } - } - - await DatabaseService.deleteMessageCascading(activeConv.id, messageId); - await conversationsStore.refreshActiveMessages(); - - conversationsStore.updateConversationTimestamp(); - } catch (error) { - console.error('Failed to delete message:', error); - } - } - - /** - * Open a fresh assistant turn anchored at the last tool result of a resolved - * agentic round and let streamChatCompletion route through runAgenticFlow. - * Used by continueAssistantMessage when classifyContinueIntent returns - * next_turn, meaning the target assistant already has its tool_calls paired - * with trailing tool results and the next thing to generate is a brand new - * turn rather than a token level continuation. - */ - private async continueAsNextAgenticTurn(anchorIndex: number): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - const anchor = conversationsStore.activeMessages[anchorIndex]; - - if (!anchor) return; - - this.cancelPreEncode(); - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const anchorMessage = findMessageById(allMessages, anchor.id); - - if (!anchorMessage) { - this.setChatLoading(activeConv.id, false); - - return; - } - - const newAssistantMessage = await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId: activeConv.id, - model: null, - role: MessageRole.ASSISTANT, - timestamp: Date.now(), - toolCalls: '', - type: MessageType.TEXT - }, - anchorMessage.id - ); - - await conversationsStore.updateCurrentNode(newAssistantMessage.id); - conversationsStore.updateConversationTimestamp(); - await conversationsStore.refreshActiveMessages(); - const conversationPath = filterByLeafNodeId( - allMessages, - anchorMessage.id, - false - ) as DatabaseMessage[]; - - await this.streamChatCompletion(conversationPath, newAssistantMessage); - } catch (error) { - if (!isAbortError(error)) console.error('Failed to continue agentic turn:', error); - - this.setChatLoading(activeConv.id, false); - } - } - - async continueAssistantMessage(messageId: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; - - const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); - - if (!result) return; - - const { index: idx, message: msg } = result; - // Decide which resume path applies. tool_calls without tool results can - // not be resumed mid sequence by continue_final_message, branch instead. - // tool_calls already paired with tool results need a fresh next turn, - // not a token level continuation of the target assistant. - const intent = classifyContinueIntent(conversationsStore.activeMessages, idx); - - if (intent.kind === ContinueIntentKind.RERUN_TURN) { - return this.regenerateMessageWithBranching(messageId); - } - - if (intent.kind === ContinueIntentKind.NEXT_TURN) { - return this.continueAsNextAgenticTurn(intent.truncateAfter); - } - - try { - this.showErrorDialog(null); - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const dbMessage = findMessageById(allMessages, messageId); - - if (!dbMessage) { - this.setChatLoading(activeConv.id, false); - - return; - } - - const originalContent = dbMessage.content; - const originalReasoning = dbMessage.reasoningContent || ''; - // Hand the persisted DatabaseMessage straight to sendMessage so its - // internal converter preserves tool_calls and extras when present. - // Reconstructing a bare {role, content} here would drop those fields - // and break continue_final_message for messages with tool calls. - const contextWithContinue = conversationsStore.activeMessages.slice(0, idx + 1); - - let appendedContent = ''; - let appendedReasoning = ''; - let hasReceivedContent = false; - - const updateStreamingContent = (fullContent: string) => { - this.setChatStreaming(msg.convId, fullContent, msg.id); - // resolve the row by id on every write, switching to another conv mid continue makes - // this a no op instead of writing positionally into the now displayed conversation - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { - content: fullContent - }); - }; - const abortController = this.getOrCreateAbortController(msg.convId); - - await ChatService.sendMessage( - contextWithContinue, - { - ...this.getApiOptions(), - continueFinalMessage: true, - onChunk: (chunk: string) => { - appendedContent += chunk; - hasReceivedContent = true; - updateStreamingContent(originalContent + appendedContent); - this.setChatReasoning(msg.convId, false); - }, - onComplete: async ( - finalContent?: string, - reasoningContent?: string, - timings?: ChatMessageTimings - ) => { - const finalAppendedContent = hasReceivedContent ? appendedContent : finalContent || ''; - const finalAppendedReasoning = hasReceivedContent - ? appendedReasoning - : reasoningContent || ''; - const fullContent = originalContent + finalAppendedContent; - const fullReasoning = originalReasoning + finalAppendedReasoning || undefined; - - await DatabaseService.updateMessage(msg.id, { - content: fullContent, - reasoningContent: fullReasoning, - timestamp: Date.now(), - timings - }); - - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { - content: fullContent, - reasoningContent: fullReasoning, - timestamp: Date.now(), - timings - }); - - conversationsStore.updateConversationTimestamp(msg.convId); - - this.setChatLoading(msg.convId, false); - this.clearChatStreaming(msg.convId); - this.setProcessingState(msg.convId, null); - }, - onCompletionId: (id: string) => { - if (!id) return; - - // refresh the message id so a later skip targets the live slot after a continue - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { - completionId: id - }); - DatabaseService.updateMessage(msg.id, { completionId: id }).catch(() => {}); - }, - onConnectionState: (state: StreamConnectionState) => { - if (msg.convId === conversationsStore.activeConversation?.id) { - this.streamConnectionState = state; - } - }, - onError: async (error: Error) => { - if (isAbortError(error)) { - if (hasReceivedContent && appendedContent) { - await DatabaseService.updateMessage(msg.id, { - content: originalContent + appendedContent, - reasoningContent: originalReasoning + appendedReasoning || undefined, - timestamp: Date.now() - }); - - conversationsStore.updateMessageAtIndex( - conversationsStore.findMessageIndex(msg.id), - { - content: originalContent + appendedContent, - reasoningContent: originalReasoning + appendedReasoning || undefined, - timestamp: Date.now() - } - ); - } - - this.setChatLoading(msg.convId, false); - this.clearChatStreaming(msg.convId); - this.setProcessingState(msg.convId, null); - - return; - } - - console.error('Continue generation error:', error); - // keep whatever was appended so far, the message stays in memory and in DB - await DatabaseService.updateMessage(msg.id, { - content: originalContent + appendedContent, - reasoningContent: originalReasoning + appendedReasoning || undefined, - timestamp: Date.now() - }); - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { - content: originalContent + appendedContent, - reasoningContent: originalReasoning + appendedReasoning || undefined, - timestamp: Date.now() - }); - - this.setChatLoading(msg.convId, false); - this.clearChatStreaming(msg.convId); - this.setProcessingState(msg.convId, null); - this.showErrorDialog({ - message: error.message, - type: error.name === 'TimeoutError' ? ErrorDialogType.TIMEOUT : ErrorDialogType.SERVER - }); - }, - onReasoningChunk: (chunk: string) => { - appendedReasoning += chunk; - hasReceivedContent = true; - // mark streaming state so a stop mid-thinking can persist the partial reasoning - this.setChatStreaming(msg.convId, originalContent + appendedContent, msg.id); - conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { - reasoningContent: originalReasoning + appendedReasoning - }); - this.setChatReasoning(msg.convId, true); - }, - onTimings: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => { - const tokensPerSecond = - timings?.predicted_ms && timings?.predicted_n - ? (timings.predicted_n / timings.predicted_ms) * 1000 - : 0; - - this.updateProcessingStateFromTimings( - { - cache_n: timings?.cache_n || 0, - predicted_n: timings?.predicted_n || 0, - predicted_per_second: tokensPerSecond, - prompt_ms: timings?.prompt_ms, - prompt_n: timings?.prompt_n || 0, - prompt_progress: promptProgress - }, - msg.convId - ); - } - }, - - msg.convId, - abortController.signal - ); - } catch (error) { - if (!isAbortError(error)) console.error('Failed to continue message:', error); - - if (activeConv) this.setChatLoading(activeConv.id, false); - } - } - - async editAssistantMessage( - messageId: string, - newContent: string, - shouldBranch: boolean - ): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; - - const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); - - if (!result) return; - - const { index: idx, message: msg } = result; - - try { - if (shouldBranch) { - const newMessage = await DatabaseService.createMessageBranch( - { - children: [], - content: newContent, - convId: msg.convId, - model: msg.model, - role: msg.role, - timestamp: Date.now(), - toolCalls: msg.toolCalls || '', - type: msg.type - }, - msg.parent! - ); - - await conversationsStore.updateCurrentNode(newMessage.id); - } else { - await DatabaseService.updateMessage(msg.id, { content: newContent }); - conversationsStore.updateMessageAtIndex(idx, { content: newContent }); - } - - conversationsStore.updateConversationTimestamp(); - - await conversationsStore.refreshActiveMessages(); - } catch (error) { - console.error('Failed to edit assistant message:', error); - } - } - - async editUserMessagePreserveResponses( - messageId: string, - newContent: string, - newExtras?: DatabaseMessageExtra[] - ): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - const result = this.getMessageByIdWithRole(messageId, MessageRole.USER); - - if (!result) return; - - const { index: idx, message: msg } = result; - - try { - const updateData: Partial = { content: newContent }; - - if (newExtras !== undefined) updateData.extra = JSON.parse(JSON.stringify(newExtras)); - - await DatabaseService.updateMessage(messageId, updateData); - - conversationsStore.updateMessageAtIndex(idx, updateData); - - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); - - if (rootMessage && msg.parent === rootMessage.id && newContent.trim()) { - await conversationsStore.updateConversationName( - activeConv.id, - generateConversationTitle( - newContent, - Boolean(settingsStore.config.titleGenerationUseFirstLine) - ) - ); - } - - conversationsStore.updateConversationTimestamp(); - } catch (error) { - console.error('Failed to edit user message:', error); - } - } - - async editMessageWithBranching( - messageId: string, - newContent: string, - newExtras?: DatabaseMessageExtra[] - ): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return; - - let result = this.getMessageByIdWithRole(messageId, MessageRole.USER); - - if (!result) result = this.getMessageByIdWithRole(messageId, MessageRole.SYSTEM); - - if (!result) return; - - const { index: idx, message: msg } = result; - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); - const isFirstUserMessage = - msg.role === MessageRole.USER && rootMessage && msg.parent === rootMessage.id; - const extrasToUse = - newExtras !== undefined - ? JSON.parse(JSON.stringify(newExtras)) - : msg.extra - ? JSON.parse(JSON.stringify(msg.extra)) - : undefined; - - let messageIdForResponse: string; - - const dbMsg = findMessageById(allMessages, msg.id); - const hasChildren = dbMsg ? dbMsg.children.length > 0 : msg.children.length > 0; - - if (!hasChildren) { - // No responses after this message — update in place instead of branching - const updates: Partial = { - content: newContent, - extra: extrasToUse, - timestamp: Date.now() - }; - - await DatabaseService.updateMessage(msg.id, updates); - conversationsStore.updateMessageAtIndex(idx, updates); - messageIdForResponse = msg.id; - } else { - // Has children — create a new branch as sibling - const parentId = msg.parent || rootMessage?.id; - - if (!parentId) return; - - const newMessage = await DatabaseService.createMessageBranch( - { - children: [], - content: newContent, - convId: msg.convId, - extra: extrasToUse, - model: msg.model, - role: msg.role, - timestamp: Date.now(), - toolCalls: msg.toolCalls || '', - type: msg.type - }, - parentId - ); - - await conversationsStore.updateCurrentNode(newMessage.id); - messageIdForResponse = newMessage.id; - } - - conversationsStore.updateConversationTimestamp(); - - if (isFirstUserMessage && newContent.trim()) - await conversationsStore.updateConversationName( - activeConv.id, - generateConversationTitle( - newContent, - Boolean(settingsStore.config.titleGenerationUseFirstLine) - ) - ); - - await conversationsStore.refreshActiveMessages(); - - if (msg.role === MessageRole.USER) - await this.generateResponseForMessage(messageIdForResponse); - } catch (error) { - console.error('Failed to edit message with branching:', error); - } - } - - private async generateResponseForMessage(userMessageId: string): Promise { - const activeConv = conversationsStore.activeConversation; - - if (!activeConv) return; - - this.showErrorDialog(null); - this.setChatLoading(activeConv.id, true); - this.clearChatStreaming(activeConv.id); - - try { - const allMessages = await conversationsStore.getConversationMessages(activeConv.id); - const conversationPath = filterByLeafNodeId( - allMessages, - userMessageId, - false - ) as DatabaseMessage[]; - const assistantMessage = await DatabaseService.createMessageBranch( - { - children: [], - content: '', - convId: activeConv.id, - model: null, - role: MessageRole.ASSISTANT, - timestamp: Date.now(), - toolCalls: '', - type: MessageType.TEXT - }, - userMessageId - ); - - conversationsStore.addMessageToActive(assistantMessage); - - await this.streamChatCompletion(conversationPath, assistantMessage); - } catch (error) { - console.error('Failed to generate response:', error); - this.setChatLoading(activeConv.id, false); - } - } - - private getContextTotal(): number | null { - const activeConvId = this.activeConversationId; - const activeState = activeConvId ? this.getProcessingState(activeConvId) : null; - - if (activeState && typeof activeState.contextTotal === 'number' && activeState.contextTotal > 0) - return activeState.contextTotal; - - if (serverStore.isRouterMode) { - const modelContextSize = modelsStore.selectedModelContextSize; - - if (typeof modelContextSize === 'number' && modelContextSize > 0) { - return modelContextSize; - } - } else { - const propsContextSize = serverStore.contextSize; - - if (typeof propsContextSize === 'number' && propsContextSize > 0) { - return propsContextSize; - } - } - - return null; - } - - updateProcessingStateFromTimings( - timingData: { - prompt_n: number; - prompt_ms?: number; - predicted_n: number; - predicted_per_second: number; - cache_n: number; - prompt_progress?: ChatMessagePromptProgress; - }, - conversationId?: string - ): void { - const processingState = this.parseTimingData(timingData); - - if (processingState === null) { - console.warn('Failed to parse timing data - skipping update'); - - return; - } - - const targetId = conversationId || this.activeConversationId; - - if (targetId) { - this.setProcessingState(targetId, processingState); - } - } - - private parseTimingData(timingData: Record): ApiProcessingState | null { - const cacheTokens = (timingData.cache_n as number) || 0, - predictedTokens = (timingData.predicted_n as number) || 0, - promptMs = (timingData.prompt_ms as number) || undefined, - promptTokens = (timingData.prompt_n as number) || 0, - tokensPerSecond = (timingData.predicted_per_second as number) || 0; - const promptProgress = timingData.prompt_progress as - | { total: number; cache: number; processed: number; time_ms: number } - | undefined; - const contextTotal = this.getContextTotal(); - const currentConfig = settingsStore.config; - const outputTokensMax = currentConfig.max_tokens || -1; - const contextUsed = promptTokens + cacheTokens + predictedTokens, - outputTokensUsed = predictedTokens; - const progressCache = promptProgress?.cache || 0, - progressActualDone = (promptProgress?.processed ?? 0) - progressCache, - progressActualTotal = (promptProgress?.total ?? 0) - progressCache; - const progressPercent = promptProgress - ? Math.round((progressActualDone / progressActualTotal) * 100) - : undefined; - - return { - cacheTokens, - contextTotal, - contextUsed, - hasNextToken: predictedTokens > 0, - outputTokensMax, - outputTokensUsed, - progressPercent, - promptMs, - promptProgress, - promptTokens, - speculative: false, - status: predictedTokens > 0 ? 'generating' : promptProgress ? 'preparing' : 'idle', - temperature: currentConfig.temperature ?? 0.8, - tokensDecoded: predictedTokens, - tokensPerSecond, - tokensRemaining: outputTokensMax - predictedTokens, - topP: currentConfig.top_p ?? 0.95 - }; - } - - restoreProcessingStateFromMessages(messages: DatabaseMessage[], conversationId: string): void { - for (let i = messages.length - 1; i >= 0; i--) { - const message = messages[i]; - - if (message.role === MessageRole.ASSISTANT && message.timings) { - const restoredState = this.parseTimingData({ - cache_n: message.timings.cache_n || 0, - predicted_n: message.timings.predicted_n || 0, - predicted_per_second: - message.timings.predicted_n && message.timings.predicted_ms - ? (message.timings.predicted_n / message.timings.predicted_ms) * 1000 - : 0, - prompt_ms: message.timings.prompt_ms, - prompt_n: message.timings.prompt_n || 0 - }); - - if (restoredState) { - this.setProcessingState(conversationId, restoredState); - - return; - } - } - } - } - - private getApiOptions(): Record { - const currentConfig = settingsStore.config; - const hasValue = (value: unknown): boolean => - value !== undefined && value !== null && value !== ''; - const apiOptions: Record = { stream: true, timings_per_token: true }; - - if (serverStore.isRouterMode) { - const modelName = modelsStore.selectedModelName; - - if (modelName) apiOptions.model = modelName; - } - - if (currentConfig.systemMessage) apiOptions.systemMessage = currentConfig.systemMessage; - - if (currentConfig.disableReasoningParsing) apiOptions.disableReasoningParsing = true; - - if (currentConfig.excludeReasoningFromContext) apiOptions.excludeReasoningFromContext = true; - - // an explicit reasoning choice overrides the server default, DEFAULT sends nothing - const effort = conversationsStore.getReasoningEffort(); - - if (effort !== ReasoningEffort.DEFAULT) { - apiOptions.enableThinking = effort !== ReasoningEffort.OFF; - - if (effort !== ReasoningEffort.OFF) apiOptions.reasoningEffort = effort; - } - - if (hasValue(currentConfig.temperature)) - apiOptions.temperature = Number(currentConfig.temperature); - - if (hasValue(currentConfig.max_tokens)) - apiOptions.max_tokens = Number(currentConfig.max_tokens); - - if (hasValue(currentConfig.dynatemp_range)) - apiOptions.dynatemp_range = Number(currentConfig.dynatemp_range); - - if (hasValue(currentConfig.dynatemp_exponent)) - apiOptions.dynatemp_exponent = Number(currentConfig.dynatemp_exponent); - - if (hasValue(currentConfig.top_k)) apiOptions.top_k = Number(currentConfig.top_k); - - if (hasValue(currentConfig.top_p)) apiOptions.top_p = Number(currentConfig.top_p); - - if (hasValue(currentConfig.min_p)) apiOptions.min_p = Number(currentConfig.min_p); - - if (hasValue(currentConfig.xtc_probability)) - apiOptions.xtc_probability = Number(currentConfig.xtc_probability); - - if (hasValue(currentConfig.xtc_threshold)) - apiOptions.xtc_threshold = Number(currentConfig.xtc_threshold); - - if (hasValue(currentConfig.typ_p)) apiOptions.typ_p = Number(currentConfig.typ_p); - - if (hasValue(currentConfig.repeat_last_n)) - apiOptions.repeat_last_n = Number(currentConfig.repeat_last_n); - - if (hasValue(currentConfig.repeat_penalty)) - apiOptions.repeat_penalty = Number(currentConfig.repeat_penalty); - - if (hasValue(currentConfig.presence_penalty)) - apiOptions.presence_penalty = Number(currentConfig.presence_penalty); - - if (hasValue(currentConfig.frequency_penalty)) - apiOptions.frequency_penalty = Number(currentConfig.frequency_penalty); - - if (hasValue(currentConfig.dry_multiplier)) - apiOptions.dry_multiplier = Number(currentConfig.dry_multiplier); - - if (hasValue(currentConfig.dry_base)) apiOptions.dry_base = Number(currentConfig.dry_base); - - if (hasValue(currentConfig.dry_allowed_length)) - apiOptions.dry_allowed_length = Number(currentConfig.dry_allowed_length); - - if (hasValue(currentConfig.dry_penalty_last_n)) - apiOptions.dry_penalty_last_n = Number(currentConfig.dry_penalty_last_n); - - if (currentConfig.samplers) apiOptions.samplers = currentConfig.samplers; - - if (hasValue(currentConfig.backend_sampling)) - apiOptions.backend_sampling = currentConfig.backend_sampling; - - if (currentConfig.customJson) apiOptions.custom = currentConfig.customJson; - - return apiOptions; - } - - private cancelPreEncode(): void { - if (this.preEncodeAbortController) { - this.preEncodeAbortController.abort(); - this.preEncodeAbortController = null; - } - } - - private async triggerPreEncode( - allMessages: DatabaseMessage[], - assistantMessage: DatabaseMessage, - assistantContent: string, - model?: string | null, - excludeReasoning?: boolean - ): Promise { - this.cancelPreEncode(); - this.preEncodeAbortController = new AbortController(); - - const signal = this.preEncodeAbortController.signal; - - try { - const allIdle = await ChatService.areAllSlotsIdle(model, signal); - - if (!allIdle || signal.aborted) return; - - const messagesWithAssistant: DatabaseMessage[] = [ - ...allMessages, - { ...assistantMessage, content: assistantContent } - ]; - - await ChatService.preEncode(messagesWithAssistant, model, excludeReasoning, signal); - } catch (err) { - if (!isAbortError(err)) { - console.warn('[ChatStore] Pre-encode failed:', err); - } - } - } -} - -export const chatStore = new ChatStore(); diff --git a/tools/ui/src/lib/stores/chat/activity.svelte.ts b/tools/ui/src/lib/stores/chat/activity.svelte.ts new file mode 100644 index 000000000000..cd4e0497bf9b --- /dev/null +++ b/tools/ui/src/lib/stores/chat/activity.svelte.ts @@ -0,0 +1,74 @@ +/** + * ChatActivityStore - Conversation activity ledger + * + * Single owner of the "is this conversation doing something" state: + * - `local` - this browser is piping a stream (send, server-stream attach, + * or resume-wait while the owning model loads) + * - `remote` - the backend reports a running session, no local pipe yet + * (global snapshot on mount / visibilitychange) + * + * The union of both drives the sidebar spinners (`loadingConvs`); `local` + * drives the per-conversation loading flags. When a local pipe ends it is + * the authoritative observer of session end, so it also drops the stale + * remote hint in the same call - no cross-owner cleanup, no ghosted + * spinners waiting for the next visibilitychange snapshot. + * + * Composed under chatStore.activity; not exported from the stores barrel. + */ + +import { SvelteSet } from 'svelte/reactivity'; + +export class ChatActivityStore { + /** Convs this browser is piping a stream for (send, attach, resume-wait). */ + private local = new SvelteSet(); + /** Convs the backend reports as having a running session (snapshot sync). */ + private remote = new SvelteSet(); + + /** Convs with any activity, the union the sidebar spinners render. */ + loadingConvs = $derived.by(() => { + const out = new SvelteSet(this.local); + + for (const id of this.remote) out.add(id); + + return Array.from(out); + }); + + /** + * Apply a backend snapshot of running sessions (mount / visibilitychange). + * Diffed so unchanged entries do not re-trigger reactivity. + */ + applyRemoteSnapshot(running: Iterable): void { + const next = new SvelteSet(running); + + for (const id of Array.from(this.remote)) { + if (!next.has(id)) this.remote.delete(id); + } + + for (const id of next) this.remote.add(id); + } + + isLocal(convId: string): boolean { + return this.local.has(convId); + } + + isRemote(convId: string): boolean { + return this.remote.has(convId); + } + + /** + * A local pipe ended for the conv. Also drops the remote hint: the local + * pipe is the authoritative observer of session end, so the sidebar hint + * goes away right away instead of ghosting until the next snapshot. + */ + localEnded(convId: string): void { + this.local.delete(convId); + this.remote.delete(convId); + } + + /** A local pipe (send, attach or resume-wait) started for the conv. */ + markLocal(convId: string): void { + this.local.add(convId); + } +} + +export const chatActivityStore = new ChatActivityStore(); diff --git a/tools/ui/src/lib/stores/context-stats.svelte.ts b/tools/ui/src/lib/stores/chat/context-stats.svelte.ts similarity index 56% rename from tools/ui/src/lib/stores/context-stats.svelte.ts rename to tools/ui/src/lib/stores/chat/context-stats.svelte.ts index 14918456300e..b5d22cfbda34 100644 --- a/tools/ui/src/lib/stores/context-stats.svelte.ts +++ b/tools/ui/src/lib/stores/chat/context-stats.svelte.ts @@ -1,5 +1,5 @@ /** - * contextStatsStore - Context window usage stats for the active conversation + * ContextStatsStore - Context window usage stats for the active conversation * * Combines token usage persisted in message timings metadata with * server-originating data: model context size from /props (modelsStore) @@ -8,12 +8,17 @@ import { MessageRole } from '$lib/enums'; // direct imports between stores, not via the barrel, to avoid circular deps -import { agenticStore } from '$lib/stores/agentic.svelte'; -import { chatStore } from '$lib/stores/chat.svelte'; -import { conversationsStore } from '$lib/stores/conversations.svelte'; -import { modelsStore } from '$lib/stores/models.svelte'; +import { agenticStore } from '$lib/stores/agentic/index.svelte'; +import { chatStore } from '$lib/stores/chat/index.svelte'; +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; import { serverStore } from '$lib/stores/server.svelte'; -import type { ApiProcessingState, ChatMessageTimings, DatabaseMessage } from '$lib/types'; +import type { + ApiProcessingState, + ChatMessageAgenticTimings, + ChatMessageTimings, + DatabaseMessage +} from '$lib/types'; interface LiveStats { freshTokens: number; @@ -22,14 +27,46 @@ interface LiveStats { outputTokens: number; } -function lastAssistantTimings(messages: DatabaseMessage[]): ChatMessageTimings | undefined { - for (let i = messages.length - 1; i >= 0; i--) { - const m = messages[i]; +interface AssistantTimingsSummary { + lastAgenticLlm: ChatMessageAgenticTimings['llm'] | undefined; + lastTimings: ChatMessageTimings | undefined; + cacheTotal: number; + output: number; + outputMs: number; + read: number; +} + +/** + * One forward pass over the messages computing everything the deriveds + * below need: the last assistant timings (per-turn gauges), the last + * agentic llm totals (cumulative gauge) and the cumulative sums. During + * streaming activeMessages churns every chunk, and each of these used to be + * its own O(n) scan re-run per chunk. + */ +function summarizeAssistantTimings(messages: DatabaseMessage[]): AssistantTimingsSummary { + let lastAgenticLlm: ChatMessageAgenticTimings['llm'] | undefined; + let lastTimings: ChatMessageTimings | undefined; + let read = 0; + let cacheTotal = 0; + let output = 0; + let outputMs = 0; + + for (const m of messages) { + if (m.role !== MessageRole.ASSISTANT || !m.timings) continue; + + lastTimings = m.timings; + + if (m.timings.agentic?.llm?.predicted_n != null) { + lastAgenticLlm = m.timings.agentic.llm; + } - if (m.role === MessageRole.ASSISTANT && m.timings) return m.timings; + read += m.timings.prompt_n ?? 0; + cacheTotal += m.timings.cache_n ?? 0; + output += m.timings.predicted_n ?? 0; + outputMs += m.timings.predicted_ms ?? 0; } - return undefined; + return { cacheTotal, lastAgenticLlm, lastTimings, output, outputMs, read }; } function deriveLiveStats(state: ApiProcessingState | null): LiveStats | null { @@ -52,25 +89,73 @@ class ContextStatsStore { // The canonical resolution lives in modelsStore.activeModelId. activeModelId = $derived(modelsStore.activeModelId); - isActiveModelLoaded = $derived( - this.activeModelId !== null && - (!serverStore.isRouterMode || modelsStore.isModelLoaded(this.activeModelId)) + // shared by currentRead/Fresh/Cache/Output and cumulative so a per-chunk + // churn of activeMessages triggers exactly one scan instead of one per + // derived + private assistantTimings = $derived.by(() => + summarizeAssistantTimings(conversationsStore.activeMessages as DatabaseMessage[]) ); - isActiveModelLoading = $derived( - this.activeModelId !== null && modelsStore.isModelOperationInProgress(this.activeModelId) - ); + private cumulative = $derived.by(() => { + const convId = conversationsStore.activeConversation?.id; + // A running agentic flow stamps llm totals on messages only when it + // exits, so read its live session totals instead. + const liveLlm = convId ? agenticStore.getLiveLlmTotals(convId) : null; + + if (liveLlm) { + const outputMs = liveLlm.predicted_ms; + const averageTokensPerSecond = + outputMs > 0 && liveLlm.predicted_n > 0 ? (liveLlm.predicted_n / outputMs) * 1000 : null; + + return { + averageTokensPerSecond, + cacheTotal: 0, + output: liveLlm.predicted_n, + read: liveLlm.prompt_n + }; + } + + const { cacheTotal, lastAgenticLlm, output, outputMs, read } = this.assistantTimings; + + // Agentic sessions stamp the same agentic.llm totals onto every + // assistant message; cache_n is never per-turn so cache_total stays 0. + if (lastAgenticLlm) { + const averageTokensPerSecond = + lastAgenticLlm.predicted_ms > 0 && lastAgenticLlm.predicted_n > 0 + ? (lastAgenticLlm.predicted_n / lastAgenticLlm.predicted_ms) * 1000 + : null; + + return { + averageTokensPerSecond, + cacheTotal: 0, + output: lastAgenticLlm.predicted_n ?? 0, + read: lastAgenticLlm.prompt_n ?? 0 + }; + } + + const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null; + + return { averageTokensPerSecond, cacheTotal, output, read }; + }); + + averageTokensPerSecond = $derived(this.cumulative.averageTokensPerSecond); contextTotal = $derived.by(() => { - void modelsStore.propsCacheVersion; + void modelsStore.props.cacheVersion; - return this.activeModelId ? modelsStore.getModelContextSize(this.activeModelId) : null; + return this.activeModelId ? modelsStore.props.getModelContextSize(this.activeModelId) : null; }); - private liveStats = $derived(deriveLiveStats(chatStore.activeProcessingState)); + private liveStats = $derived(deriveLiveStats(chatStore.processing.activeState)); + + currentOutput = $derived.by(() => { + if (this.liveStats && this.liveStats.outputTokens > 0) return this.liveStats.outputTokens; + + return this.assistantTimings.lastTimings?.predicted_n ?? 0; + }); currentRead = $derived.by(() => { - const timings = lastAssistantTimings(conversationsStore.activeMessages as DatabaseMessage[]); + const timings = this.assistantTimings.lastTimings; let read = 0; @@ -87,34 +172,6 @@ class ContextStatsStore { return read; }); - currentFresh = $derived.by(() => { - const timings = lastAssistantTimings(conversationsStore.activeMessages as DatabaseMessage[]); - const fresh = timings?.prompt_n ?? 0; - - return Math.max(fresh, this.liveStats?.freshTokens ?? 0); - }); - - currentCache = $derived.by(() => { - const timings = lastAssistantTimings(conversationsStore.activeMessages as DatabaseMessage[]); - const cached = timings?.cache_n ?? 0; - - if (this.liveStats && this.liveStats.promptTokens > 0) { - return Math.max(cached, this.liveStats.cacheTokens); - } - - return cached; - }); - - currentOutput = $derived.by(() => { - if (this.liveStats && this.liveStats.outputTokens > 0) return this.liveStats.outputTokens; - - const timings = lastAssistantTimings(conversationsStore.activeMessages as DatabaseMessage[]); - - return timings?.predicted_n ?? 0; - }); - - kvTotal = $derived(this.currentRead + this.currentOutput); - contextUsed = $derived(this.currentRead + this.currentOutput); contextAvailable = $derived( @@ -127,71 +184,38 @@ class ContextStatsStore { return Math.round((this.contextUsed / this.contextTotal) * 100); }); - private cumulative = $derived.by(() => { - const messages = conversationsStore.activeMessages as DatabaseMessage[]; - const convId = conversationsStore.activeConversation?.id; - // A running agentic flow stamps llm totals on messages only when it - // exits, so read its live session totals instead. - const liveLlm = convId ? agenticStore.getLiveLlmTotals(convId) : null; - - if (liveLlm) { - const outputMs = liveLlm.predicted_ms; - const averageTokensPerSecond = - outputMs > 0 && liveLlm.predicted_n > 0 ? (liveLlm.predicted_n / outputMs) * 1000 : null; + cumulativeCacheTotal = $derived(this.cumulative.cacheTotal); - return { - averageTokensPerSecond, - cacheTotal: 0, - output: liveLlm.predicted_n, - read: liveLlm.prompt_n - }; - } + cumulativeOutput = $derived(this.cumulative.output); - // Agentic sessions stamp the same agentic.llm totals onto every - // assistant message; cache_n is never per-turn so cache_total stays 0. - const agenticMessages = messages.filter( - (m) => m.role === MessageRole.ASSISTANT && m.timings?.agentic?.llm?.predicted_n != null - ); + cumulativeRead = $derived(this.cumulative.read); - if (agenticMessages.length > 0) { - const llm = agenticMessages[agenticMessages.length - 1].timings!.agentic!.llm; - const output = llm.predicted_n ?? 0; - const outputMs = llm.predicted_ms ?? 0; - const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null; + currentCache = $derived.by(() => { + const cached = this.assistantTimings.lastTimings?.cache_n ?? 0; - return { - averageTokensPerSecond, - cacheTotal: 0, - output, - read: llm.prompt_n ?? 0 - }; + if (this.liveStats && this.liveStats.promptTokens > 0) { + return Math.max(cached, this.liveStats.cacheTokens); } - let read = 0; - let output = 0; - let outputMs = 0; - let cacheTotal = 0; - - for (const m of messages) { - if (m.role !== MessageRole.ASSISTANT || !m.timings) continue; + return cached; + }); - read += m.timings.prompt_n ?? 0; - cacheTotal += m.timings.cache_n ?? 0; - output += m.timings.predicted_n ?? 0; - outputMs += m.timings.predicted_ms ?? 0; - } - const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null; + currentFresh = $derived.by(() => { + const fresh = this.assistantTimings.lastTimings?.prompt_n ?? 0; - return { averageTokensPerSecond, cacheTotal, output, read }; + return Math.max(fresh, this.liveStats?.freshTokens ?? 0); }); - cumulativeRead = $derived(this.cumulative.read); - - cumulativeOutput = $derived(this.cumulative.output); + isActiveModelLoaded = $derived( + this.activeModelId !== null && + (!serverStore.isRouterMode || modelsStore.isModelLoaded(this.activeModelId)) + ); - cumulativeCacheTotal = $derived(this.cumulative.cacheTotal); + isActiveModelLoading = $derived( + this.activeModelId !== null && modelsStore.status.isOperationInProgress(this.activeModelId) + ); - averageTokensPerSecond = $derived(this.cumulative.averageTokensPerSecond); + kvTotal = $derived(this.currentRead + this.currentOutput); } export const contextStatsStore = new ContextStatsStore(); diff --git a/tools/ui/src/lib/stores/draft-messages.svelte.ts b/tools/ui/src/lib/stores/chat/drafts.svelte.ts similarity index 76% rename from tools/ui/src/lib/stores/draft-messages.svelte.ts rename to tools/ui/src/lib/stores/chat/drafts.svelte.ts index 235a59122e7f..f480e1efd490 100644 --- a/tools/ui/src/lib/stores/draft-messages.svelte.ts +++ b/tools/ui/src/lib/stores/chat/drafts.svelte.ts @@ -1,3 +1,11 @@ +/** + * DraftMessagesStore - Per-conversation input drafts + * + * Keeps in-memory drafts (message text + files) keyed by conversation id, + * plus a dedicated key for the new-chat screen, so the input box restores + * its content when switching conversations. + */ + import { NEW_CHAT_DRAFT_KEY } from '$lib/constants'; interface DraftMessage { @@ -8,6 +16,12 @@ interface DraftMessage { class DraftMessagesStore { private drafts = new Map(); + clearDraftMessage(chatId: string | undefined): void { + const key = chatId ?? NEW_CHAT_DRAFT_KEY; + + this.drafts.delete(key); + } + getDraftMessage(chatId: string | undefined): DraftMessage { const key = chatId ?? NEW_CHAT_DRAFT_KEY; @@ -23,12 +37,6 @@ class DraftMessagesStore { this.drafts.delete(key); } } - - clearDraftMessage(chatId: string | undefined): void { - const key = chatId ?? NEW_CHAT_DRAFT_KEY; - - this.drafts.delete(key); - } } export const draftMessagesStore = new DraftMessagesStore(); diff --git a/tools/ui/src/lib/stores/chat/flows.svelte.ts b/tools/ui/src/lib/stores/chat/flows.svelte.ts new file mode 100644 index 000000000000..16c377bb6127 --- /dev/null +++ b/tools/ui/src/lib/stores/chat/flows.svelte.ts @@ -0,0 +1,794 @@ +/** + * ChatMessageFlows - Message-level flows for the active conversation + * + * Owns the operations that mutate chat history and (re)stream a response: + * editing, regeneration, continuation and deletion of messages. Created and + * owned by chatStore; the host exposes the streaming core and the + * per-conversation state setters these flows drive. + */ + +import { + ContinueIntentKind, + ErrorDialogType, + MessageRole, + MessageType, + StreamConnectionState +} from '$lib/enums'; +import { ChatService } from '$lib/services/chat.service'; +import { DatabaseService } from '$lib/services/database.service'; +import type { ChatProcessingStore } from '$lib/stores/chat/processing.svelte'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import type { + ChatMessagePromptProgress, + ChatMessageTimings, + DatabaseMessage, + DatabaseMessageExtra, + ErrorDialogState +} from '$lib/types'; +import { + classifyContinueIntent, + filterByLeafNodeId, + findDescendantMessages, + findLeafNode, + findMessageById, + isAbortError +} from '$lib/utils'; + +/** + * The slice of chatStore the flows drive. Kept narrow on purpose so the flows + * cannot reach around the host's full surface; chatStore implements this + * structurally. + */ +export interface ChatFlowsHost { + processing: ChatProcessingStore; + streamConnectionState: StreamConnectionState; + cancelPreEncode(): void; + clearChatStreaming(convId: string, messageId?: string): void; + cleanupStreaming(convId: string): void; + createAssistantMessage(parentId?: string): Promise; + getApiOptions(): Record; + getOrCreateAbortController(convId: string): AbortController; + isChatLoadingInternal(convId: string): boolean; + setChatLoading(convId: string, loading: boolean): void; + setChatReasoning(convId: string, reasoning: boolean): void; + setChatStreaming( + convId: string, + response: string, + messageId: string, + model?: string | null + ): void; + showErrorDialog(state: ErrorDialogState | null): void; + stopGeneration(): Promise; + streamChatCompletion( + allMessages: DatabaseMessage[], + assistantMessage: DatabaseMessage, + onComplete?: (content: string) => Promise, + onError?: (error: Error) => void, + modelOverride?: string | null, + firstUserMessageContent?: string + ): Promise; +} + +export class ChatMessageFlows { + constructor(private host: ChatFlowsHost) {} + + async continueAssistantMessage(messageId: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv || this.host.isChatLoadingInternal(activeConv.id)) return; + + const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); + + if (!result) return; + + const { index: idx, message: msg } = result; + // Decide which resume path applies. tool_calls without tool results can + // not be resumed mid sequence by continue_final_message, branch instead. + // tool_calls already paired with tool results need a fresh next turn, + // not a token level continuation of the target assistant. + const intent = classifyContinueIntent(conversationsStore.activeMessages, idx); + + if (intent.kind === ContinueIntentKind.RERUN_TURN) { + return this.regenerateMessageWithBranching(messageId); + } + + if (intent.kind === ContinueIntentKind.NEXT_TURN) { + return this.continueAsNextAgenticTurn(intent.truncateAfter); + } + + try { + this.host.showErrorDialog(null); + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const dbMessage = findMessageById(allMessages, messageId); + + if (!dbMessage) { + this.host.setChatLoading(activeConv.id, false); + + return; + } + + const originalContent = dbMessage.content; + const originalReasoning = dbMessage.reasoningContent || ''; + // Hand the persisted DatabaseMessage straight to sendMessage so its + // internal converter preserves tool_calls and extras when present. + // Reconstructing a bare {role, content} here would drop those fields + // and break continue_final_message for messages with tool calls. + const contextWithContinue = conversationsStore.activeMessages.slice(0, idx + 1); + + let appendedContent = ''; + let appendedReasoning = ''; + let hasReceivedContent = false; + + const updateStreamingContent = (fullContent: string) => { + this.host.setChatStreaming(msg.convId, fullContent, msg.id); + // resolve the row by id on every write, switching to another conv mid continue makes + // this a no op instead of writing positionally into the now displayed conversation + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + content: fullContent + }); + }; + const abortController = this.host.getOrCreateAbortController(msg.convId); + + await ChatService.sendMessage( + contextWithContinue, + { + ...this.host.getApiOptions(), + continueFinalMessage: true, + onChunk: (chunk: string) => { + appendedContent += chunk; + hasReceivedContent = true; + updateStreamingContent(originalContent + appendedContent); + this.host.setChatReasoning(msg.convId, false); + }, + onComplete: async ( + finalContent?: string, + reasoningContent?: string, + timings?: ChatMessageTimings + ) => { + const finalAppendedContent = hasReceivedContent ? appendedContent : finalContent || ''; + const finalAppendedReasoning = hasReceivedContent + ? appendedReasoning + : reasoningContent || ''; + const fullContent = originalContent + finalAppendedContent; + const fullReasoning = originalReasoning + finalAppendedReasoning || undefined; + + await DatabaseService.updateMessage(msg.id, { + content: fullContent, + reasoningContent: fullReasoning, + timestamp: Date.now(), + timings + }); + + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + content: fullContent, + reasoningContent: fullReasoning, + timestamp: Date.now(), + timings + }); + + conversationsStore.updateConversationTimestamp(msg.convId); + + this.host.cleanupStreaming(msg.convId); + }, + onCompletionId: (id: string) => { + if (!id) return; + + // refresh the message id so a later skip targets the live slot after a continue + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + completionId: id + }); + DatabaseService.updateMessage(msg.id, { completionId: id }).catch(() => {}); + }, + onConnectionState: (state: StreamConnectionState) => { + if (msg.convId === conversationsStore.activeConversation?.id) { + this.host.streamConnectionState = state; + } + }, + onError: async (error: Error) => { + if (isAbortError(error)) { + if (hasReceivedContent && appendedContent) { + await DatabaseService.updateMessage(msg.id, { + content: originalContent + appendedContent, + reasoningContent: originalReasoning + appendedReasoning || undefined, + timestamp: Date.now() + }); + + conversationsStore.updateMessageAtIndex( + conversationsStore.findMessageIndex(msg.id), + { + content: originalContent + appendedContent, + reasoningContent: originalReasoning + appendedReasoning || undefined, + timestamp: Date.now() + } + ); + } + + this.host.cleanupStreaming(msg.convId); + + return; + } + + console.error('Continue generation error:', error); + // keep whatever was appended so far, the message stays in memory and in DB + await DatabaseService.updateMessage(msg.id, { + content: originalContent + appendedContent, + reasoningContent: originalReasoning + appendedReasoning || undefined, + timestamp: Date.now() + }); + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + content: originalContent + appendedContent, + reasoningContent: originalReasoning + appendedReasoning || undefined, + timestamp: Date.now() + }); + + this.host.cleanupStreaming(msg.convId); + this.host.showErrorDialog({ + message: error.message, + type: error.name === 'TimeoutError' ? ErrorDialogType.TIMEOUT : ErrorDialogType.SERVER + }); + }, + onReasoningChunk: (chunk: string) => { + appendedReasoning += chunk; + hasReceivedContent = true; + // mark streaming state so a stop mid-thinking can persist the partial reasoning + this.host.setChatStreaming(msg.convId, originalContent + appendedContent, msg.id); + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + reasoningContent: originalReasoning + appendedReasoning + }); + this.host.setChatReasoning(msg.convId, true); + }, + onTimings: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => { + this.host.processing.applyStreamTimings(timings, promptProgress, msg.convId); + } + }, + + msg.convId, + abortController.signal + ); + } catch (error) { + if (!isAbortError(error)) console.error('Failed to continue message:', error); + + if (activeConv) this.host.setChatLoading(activeConv.id, false); + } + } + + async deleteMessage(messageId: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const messageToDelete = findMessageById(allMessages, messageId); + + if (!messageToDelete) return; + + const currentPath = filterByLeafNodeId(allMessages, activeConv.currNode || '', false); + const isInCurrentPath = currentPath.some((m) => m.id === messageId); + + if (isInCurrentPath && messageToDelete.parent) { + const siblings = allMessages.filter( + (m) => m.parent === messageToDelete.parent && m.id !== messageId + ); + + if (siblings.length > 0) { + const latestSibling = siblings.reduce((latest, sibling) => + sibling.timestamp > latest.timestamp ? sibling : latest + ); + + await conversationsStore.updateCurrentNode(findLeafNode(allMessages, latestSibling.id)); + } else if (messageToDelete.parent) { + await conversationsStore.updateCurrentNode( + findLeafNode(allMessages, messageToDelete.parent) + ); + } + } + + await DatabaseService.deleteMessageCascading(activeConv.id, messageId); + await conversationsStore.refreshActiveMessages(); + + conversationsStore.updateConversationTimestamp(); + } catch (error) { + console.error('Failed to delete message:', error); + } + } + + async editAssistantMessage( + messageId: string, + newContent: string, + shouldBranch: boolean + ): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv || this.host.isChatLoadingInternal(activeConv.id)) return; + + const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); + + if (!result) return; + + const { index: idx, message: msg } = result; + + try { + if (shouldBranch) { + const newMessage = await DatabaseService.createMessageBranch( + { + children: [], + content: newContent, + convId: msg.convId, + model: msg.model, + role: msg.role, + timestamp: Date.now(), + toolCalls: msg.toolCalls || '', + type: msg.type + }, + msg.parent! + ); + + await conversationsStore.updateCurrentNode(newMessage.id); + } else { + await DatabaseService.updateMessage(msg.id, { content: newContent }); + conversationsStore.updateMessageAtIndex(idx, { content: newContent }); + } + + conversationsStore.updateConversationTimestamp(); + + await conversationsStore.refreshActiveMessages(); + } catch (error) { + console.error('Failed to edit assistant message:', error); + } + } + + async editMessageWithBranching( + messageId: string, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv || this.host.isChatLoadingInternal(activeConv.id)) return; + + let result = this.getMessageByIdWithRole(messageId, MessageRole.USER); + + if (!result) result = this.getMessageByIdWithRole(messageId, MessageRole.SYSTEM); + + if (!result) return; + + const { index: idx, message: msg } = result; + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); + const isFirstUserMessage = + msg.role === MessageRole.USER && rootMessage && msg.parent === rootMessage.id; + const extrasToUse = + newExtras !== undefined + ? JSON.parse(JSON.stringify(newExtras)) + : msg.extra + ? JSON.parse(JSON.stringify(msg.extra)) + : undefined; + + let messageIdForResponse: string; + + const dbMsg = findMessageById(allMessages, msg.id); + const hasChildren = dbMsg ? dbMsg.children.length > 0 : msg.children.length > 0; + + if (!hasChildren) { + // No responses after this message - update in place instead of branching + const updates: Partial = { + content: newContent, + extra: extrasToUse, + timestamp: Date.now() + }; + + await DatabaseService.updateMessage(msg.id, updates); + conversationsStore.updateMessageAtIndex(idx, updates); + messageIdForResponse = msg.id; + } else { + // Has children - create a new branch as sibling + const parentId = msg.parent || rootMessage?.id; + + if (!parentId) return; + + const newMessage = await DatabaseService.createMessageBranch( + { + children: [], + content: newContent, + convId: msg.convId, + extra: extrasToUse, + model: msg.model, + role: msg.role, + timestamp: Date.now(), + toolCalls: msg.toolCalls || '', + type: msg.type + }, + parentId + ); + + await conversationsStore.updateCurrentNode(newMessage.id); + messageIdForResponse = newMessage.id; + } + + conversationsStore.updateConversationTimestamp(); + + if (isFirstUserMessage && newContent.trim()) + await conversationsStore.applyTitleFromContent(activeConv.id, newContent); + + await conversationsStore.refreshActiveMessages(); + + if (msg.role === MessageRole.USER) + await this.generateResponseForMessage(messageIdForResponse); + } catch (error) { + console.error('Failed to edit message with branching:', error); + } + } + + async editUserMessagePreserveResponses( + messageId: string, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + const result = this.getMessageByIdWithRole(messageId, MessageRole.USER); + + if (!result) return; + + const { index: idx, message: msg } = result; + + try { + const updateData: Partial = { content: newContent }; + + if (newExtras !== undefined) updateData.extra = JSON.parse(JSON.stringify(newExtras)); + + await DatabaseService.updateMessage(messageId, updateData); + + conversationsStore.updateMessageAtIndex(idx, updateData); + + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); + + if (rootMessage && msg.parent === rootMessage.id && newContent.trim()) { + await conversationsStore.applyTitleFromContent(activeConv.id, newContent); + } + + conversationsStore.updateConversationTimestamp(); + } catch (error) { + console.error('Failed to edit user message:', error); + } + } + + async getDeletionInfo(messageId: string): Promise<{ + totalCount: number; + userMessages: number; + assistantMessages: number; + messageTypes: string[]; + }> { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) + return { assistantMessages: 0, messageTypes: [], totalCount: 0, userMessages: 0 }; + + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const messageToDelete = findMessageById(allMessages, messageId); + + // For system messages, don't count descendants as they will be preserved (reparented to root) + if (messageToDelete?.role === MessageRole.SYSTEM) { + const messagesToDelete = allMessages.filter((m) => m.id === messageId); + + let assistantMessages = 0, + userMessages = 0; + + const messageTypes: string[] = []; + + for (const msg of messagesToDelete) { + if (msg.role === MessageRole.USER) { + userMessages++; + + if (!messageTypes.includes('user message')) messageTypes.push('user message'); + } else if (msg.role === MessageRole.ASSISTANT) { + assistantMessages++; + + if (!messageTypes.includes('assistant response')) messageTypes.push('assistant response'); + } + } + + return { assistantMessages, messageTypes, totalCount: 1, userMessages }; + } + + const descendants = findDescendantMessages(allMessages, messageId); + const allToDelete = [messageId, ...descendants]; + const messagesToDelete = allMessages.filter((m) => allToDelete.includes(m.id)); + + let assistantMessages = 0, + userMessages = 0; + + const messageTypes: string[] = []; + + for (const msg of messagesToDelete) { + if (msg.role === MessageRole.USER) { + userMessages++; + + if (!messageTypes.includes('user message')) messageTypes.push('user message'); + } else if (msg.role === MessageRole.ASSISTANT) { + assistantMessages++; + + if (!messageTypes.includes('assistant response')) messageTypes.push('assistant response'); + } + } + + return { assistantMessages, messageTypes, totalCount: allToDelete.length, userMessages }; + } + + async regenerateMessage(messageId: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv || this.host.isChatLoadingInternal(activeConv.id)) return; + + this.host.cancelPreEncode(); + const result = this.getMessageByIdWithRole(messageId, MessageRole.ASSISTANT); + + if (!result) return; + + const { index: messageIndex } = result; + + try { + const messagesToRemove = conversationsStore.activeMessages.slice(messageIndex); + + await DatabaseService.deleteMessageCascading(activeConv.id, messagesToRemove[0].id); + conversationsStore.sliceActiveMessages(messageIndex); + conversationsStore.updateConversationTimestamp(); + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + const parentMessageId = + conversationsStore.activeMessages.length > 0 + ? conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1].id + : undefined; + const assistantMessage = await this.host.createAssistantMessage(parentMessageId); + + conversationsStore.addMessageToActive(assistantMessage); + await this.host.streamChatCompletion( + conversationsStore.activeMessages.slice(0, -1), + assistantMessage + ); + } catch (error) { + if (!isAbortError(error)) console.error('Failed to regenerate message:', error); + + this.host.setChatLoading(activeConv?.id || '', false); + } + } + + async regenerateMessageWithBranching(messageId: string, modelOverride?: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv || this.host.isChatLoadingInternal(activeConv.id)) return; + + this.host.cancelPreEncode(); + try { + const idx = conversationsStore.findMessageIndex(messageId); + + if (idx === -1) return; + + const msg = conversationsStore.activeMessages[idx]; + + if (msg.role !== MessageRole.ASSISTANT) return; + + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const parentMessage = findMessageById(allMessages, msg.parent); + + if (!parentMessage) return; + + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + const newAssistantMessage = await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId: msg.convId, + model: null, + role: msg.role, + timestamp: Date.now(), + toolCalls: '', + type: msg.type + }, + parentMessage.id + ); + + await conversationsStore.updateCurrentNode(newAssistantMessage.id); + conversationsStore.updateConversationTimestamp(); + await conversationsStore.refreshActiveMessages(); + const conversationPath = filterByLeafNodeId( + allMessages, + parentMessage.id, + false + ) as DatabaseMessage[]; + const modelToUse = modelOverride || msg.model || undefined; + + await this.host.streamChatCompletion( + conversationPath, + newAssistantMessage, + undefined, + undefined, + modelToUse + ); + } catch (error) { + if (!isAbortError(error)) + console.error('Failed to regenerate message with branching:', error); + + this.host.setChatLoading(activeConv?.id || '', false); + } + } + + async updateMessage(messageId: string, newContent: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + if (this.host.isChatLoadingInternal(activeConv.id)) await this.host.stopGeneration(); + + const result = this.getMessageByIdWithRole(messageId, MessageRole.USER); + + if (!result) return; + + const { index: messageIndex, message: messageToUpdate } = result; + const originalContent = messageToUpdate.content; + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); + const isFirstUserMessage = rootMessage && messageToUpdate.parent === rootMessage.id; + + conversationsStore.updateMessageAtIndex(messageIndex, { content: newContent }); + await DatabaseService.updateMessage(messageId, { content: newContent }); + + if (isFirstUserMessage && newContent.trim()) + await conversationsStore.applyTitleFromContent(activeConv.id, newContent); + + const messagesToRemove = conversationsStore.activeMessages.slice(messageIndex + 1); + + if (messagesToRemove.length > 0) + await DatabaseService.deleteMessageCascading(activeConv.id, messagesToRemove[0].id); + + conversationsStore.sliceActiveMessages(messageIndex + 1); + conversationsStore.updateConversationTimestamp(); + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + const assistantMessage = await this.host.createAssistantMessage(); + + conversationsStore.addMessageToActive(assistantMessage); + await conversationsStore.updateCurrentNode(assistantMessage.id); + await this.host.streamChatCompletion( + conversationsStore.activeMessages.slice(0, -1), + assistantMessage, + undefined, + () => { + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(messageId), { + content: originalContent + }); + } + ); + } catch (error) { + if (!isAbortError(error)) console.error('Failed to update message:', error); + } + } + + /** + * Open a fresh assistant turn anchored at the last tool result of a resolved + * agentic round and let streamChatCompletion route through runAgenticFlow. + * Used by continueAssistantMessage when classifyContinueIntent returns + * next_turn, meaning the target assistant already has its tool_calls paired + * with trailing tool results and the next thing to generate is a brand new + * turn rather than a token level continuation. + */ + private async continueAsNextAgenticTurn(anchorIndex: number): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + const anchor = conversationsStore.activeMessages[anchorIndex]; + + if (!anchor) return; + + this.host.cancelPreEncode(); + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const anchorMessage = findMessageById(allMessages, anchor.id); + + if (!anchorMessage) { + this.host.setChatLoading(activeConv.id, false); + + return; + } + + const newAssistantMessage = await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId: activeConv.id, + model: null, + role: MessageRole.ASSISTANT, + timestamp: Date.now(), + toolCalls: '', + type: MessageType.TEXT + }, + anchorMessage.id + ); + + await conversationsStore.updateCurrentNode(newAssistantMessage.id); + conversationsStore.updateConversationTimestamp(); + await conversationsStore.refreshActiveMessages(); + const conversationPath = filterByLeafNodeId( + allMessages, + anchorMessage.id, + false + ) as DatabaseMessage[]; + + await this.host.streamChatCompletion(conversationPath, newAssistantMessage); + } catch (error) { + if (!isAbortError(error)) console.error('Failed to continue agentic turn:', error); + + this.host.setChatLoading(activeConv.id, false); + } + } + + private async generateResponseForMessage(userMessageId: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + this.host.showErrorDialog(null); + this.host.setChatLoading(activeConv.id, true); + this.host.clearChatStreaming(activeConv.id); + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const conversationPath = filterByLeafNodeId( + allMessages, + userMessageId, + false + ) as DatabaseMessage[]; + const assistantMessage = await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId: activeConv.id, + model: null, + role: MessageRole.ASSISTANT, + timestamp: Date.now(), + toolCalls: '', + type: MessageType.TEXT + }, + userMessageId + ); + + conversationsStore.addMessageToActive(assistantMessage); + + await this.host.streamChatCompletion(conversationPath, assistantMessage); + } catch (error) { + console.error('Failed to generate response:', error); + this.host.setChatLoading(activeConv.id, false); + } + } + + private getMessageByIdWithRole( + messageId: string, + expectedRole?: MessageRole + ): { message: DatabaseMessage; index: number } | null { + const index = conversationsStore.findMessageIndex(messageId); + + if (index === -1) return null; + + const message = conversationsStore.activeMessages[index]; + + if (expectedRole && message.role !== expectedRole) return null; + + return { index, message }; + } +} diff --git a/tools/ui/src/lib/stores/chat/index.svelte.ts b/tools/ui/src/lib/stores/chat/index.svelte.ts new file mode 100644 index 000000000000..aab824fd715e --- /dev/null +++ b/tools/ui/src/lib/stores/chat/index.svelte.ts @@ -0,0 +1,1441 @@ +/** + * chatStore - Chat lifecycle, streaming and message operations + * + * Owns the active conversation's chat state: sending messages, streaming + * responses, editing/regeneration flows and per-conversation processing + * activity. Composes the stream manager, message flows, activity ledger and + * processing snapshot; persists through conversationsStore. + * + * Uses ChatService for the API layer and conversationsStore for persistence. + */ + +import { CWD_CLEARED_TEXT, SYSTEM_MESSAGE_PLACEHOLDER, TITLE_GENERATION } from '$lib/constants'; +import { + ErrorDialogType, + MessageRole, + MessageType, + ReasoningEffort, + StreamConnectionState +} from '$lib/enums'; +import { ChatService } from '$lib/services/chat.service'; +import { DatabaseService } from '$lib/services/database.service'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { agenticStore } from '$lib/stores/agentic/index.svelte'; +import { chatActivityStore } from '$lib/stores/chat/activity.svelte'; +import { type ChatFlowsHost, ChatMessageFlows } from '$lib/stores/chat/flows.svelte'; +import { chatProcessingStore } from '$lib/stores/chat/processing.svelte'; +import { type ChatStreamHost, ChatStreamManager } from '$lib/stores/chat/streams.svelte'; +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import { mcpStore } from '$lib/stores/mcp/index.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { serverStore } from '$lib/stores/server.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; +import { toolsStore } from '$lib/stores/tools.svelte'; +import type { + ApiChatMessageData, + ChatMessagePromptProgress, + ChatMessageTimings, + ChatStreamCallbacks, + DatabaseMessage, + DatabaseMessageExtra, + ErrorDialogState +} from '$lib/types'; +import { + findMessageById, + formatCwdMessage, + getConversationModel, + isAbortError, + normalizeModelName +} from '$lib/utils'; +import { SvelteMap } from 'svelte/reactivity'; + +class ChatStore implements ChatStreamHost, ChatFlowsHost { + chatReasoningStates = new SvelteMap(); + chatStreamingStates = new SvelteMap< + string, + { response: string; messageId: string; model?: string | null } + >(); + currentResponse = $state(''); + errorDialogState = $state(null); + // true while the active conversation has a local pipe (send, attach or resume-wait) + isLoading = $derived(this.activity.isLocal(conversationsStore.activeConversation?.id ?? '')); + // true while the active conversation streams reasoning content but no visible content yet + isReasoning = $derived( + this.chatReasoningStates.get(conversationsStore.activeConversation?.id ?? '') ?? false + ); + pendingEditMessageId = $state(null); + // resumable stream connection state for the active conversation + // streaming -> bytes flowing normally, resuming -> waiting on /v1/stream reconnect, lost -> unrecoverable + streamConnectionState = $state(StreamConnectionState.STREAMING); + private abortControllers = new SvelteMap(); + private addFilesHandler: ((files: File[]) => void) | null = $state(null); + // message flows: edit, regenerate, continue, delete + private flows = new ChatMessageFlows(this); + private isEditModeActive = $state(false); + private pendingDraftFiles = $state([]); + private pendingDraftMessage = $state(''); + /** Reactive: queued pending messages for non-agentic streaming */ + private pendingMessages = new SvelteMap< + string, + { content: string; extras?: DatabaseMessageExtra[] } + >(); + private preEncodeAbortController: AbortController | null = null; + + // server-side stream sessions: discovery, attach/replay, resume retry, remote sync + private streams = new ChatStreamManager(this); + + /** Conv activity (local pipe / remote session), composed here. */ + get activity() { + return chatActivityStore; + } + + /** Processing state, composed here so consumers have a single chat scope. */ + get processing() { + return chatProcessingStore; + } + + /** + * Abort the current agentic flow signal without clearing loading state. + * Used by "Send immediately" to force the agentic loop to exit so that + * the pending steering message can be re-sent. + * + * Any tool calls captured mid-stream are dropped before the abort so the + * pending message (or a manual follow-up) does not re-send a half-received + * tool call with invalid JSON arguments to the server. Mirrors what the + * Stop button already does through stopGenerationForChat. + */ + async abortCurrentFlow(convId: string): Promise { + await this.savePartialResponseIfNeeded(convId); + const c = this.abortControllers.get(convId); + + if (c) { + c.abort(); + this.abortControllers.delete(convId); + } + } + + async addMessage( + role: MessageRole, + content: string, + type: MessageType = MessageType.TEXT, + parent: string = '-1', + extras?: DatabaseMessageExtra[], + isSynthetic?: boolean + ): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) throw new Error('No active conversation'); + + let parentId: string | null = null; + + if (parent === '-1') { + const am = conversationsStore.activeMessages; + + if (am.length > 0) parentId = am[am.length - 1].id; + else { + const all = await conversationsStore.getConversationMessages(activeConv.id); + const r = all.find((m) => m.parent === null && m.type === 'root'); + + parentId = r ? r.id : await DatabaseService.createRootMessage(activeConv.id); + } + } else parentId = parent; + + const message = await DatabaseService.createMessageBranch( + { + children: [], + content, + convId: activeConv.id, + extra: extras, + isSynthetic, + role, + timestamp: Date.now(), + toolCalls: '', + type + }, + parentId + ); + + conversationsStore.addMessageToActive(message); + await conversationsStore.updateCurrentNode(message.id); + conversationsStore.updateConversationTimestamp(); + + return message; + } + async addSystemPrompt(): Promise { + let activeConv = conversationsStore.activeConversation; + + if (!activeConv) { + await conversationsStore.createConversation(); + activeConv = conversationsStore.activeConversation; + } + + if (!activeConv) return; + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); + const rootId = rootMessage + ? rootMessage.id + : await DatabaseService.createRootMessage(activeConv.id); + const existingSystemMessage = allMessages.find( + (m) => m.role === MessageRole.SYSTEM && m.parent === rootId + ); + + if (existingSystemMessage) { + this.pendingEditMessageId = existingSystemMessage.id; + + if (!conversationsStore.activeMessages.some((m) => m.id === existingSystemMessage.id)) + conversationsStore.activeMessages.unshift(existingSystemMessage); + + return; + } + + const am = conversationsStore.activeMessages; + const firstActiveMessage = am.find((m) => m.parent === rootId); + const systemMessage = await DatabaseService.createSystemMessage( + activeConv.id, + SYSTEM_MESSAGE_PLACEHOLDER, + rootId + ); + + if (firstActiveMessage) { + await DatabaseService.updateMessage(firstActiveMessage.id, { + parent: systemMessage.id + }); + await DatabaseService.updateMessage(systemMessage.id, { + children: [firstActiveMessage.id] + }); + const updatedRootChildren = rootMessage + ? rootMessage.children.filter((id: string) => id !== firstActiveMessage.id) + : []; + + await DatabaseService.updateMessage(rootId, { + children: [ + ...updatedRootChildren.filter((id: string) => id !== systemMessage.id), + systemMessage.id + ] + }); + const firstMsgIndex = conversationsStore.findMessageIndex(firstActiveMessage.id); + + if (firstMsgIndex !== -1) + conversationsStore.updateMessageAtIndex(firstMsgIndex, { + parent: systemMessage.id + }); + } + + conversationsStore.activeMessages.unshift(systemMessage); + this.pendingEditMessageId = systemMessage.id; + conversationsStore.updateConversationTimestamp(); + } catch (error) { + console.error('Failed to add system prompt:', error); + } + } + cancelPreEncode(): void { + if (this.preEncodeAbortController) { + this.preEncodeAbortController.abort(); + this.preEncodeAbortController = null; + } + } + + /** + * Resets the loading, streaming and processing state for a conversation + * after a generation ends or errors. Shared by the flows' exit paths. + */ + cleanupStreaming(convId: string): void { + this.setChatLoading(convId, false); + this.clearChatStreaming(convId); + this.processing.setState(convId, null); + } + clearChatStreaming(convId: string, messageId?: string): void { + // session aware: a stale generation must not wipe a newer one's streaming state on the + // same conversation, that would drop the frozen stop identity and stop the wrong session + if (messageId !== undefined) { + const cur = this.chatStreamingStates.get(convId); + + if (cur && cur.messageId !== messageId) return; + } + + this.chatStreamingStates.delete(convId); + + if (convId === conversationsStore.activeConversation?.id) this.currentResponse = ''; + } + clearEditMode(): void { + this.isEditModeActive = false; + this.addFilesHandler = null; + } + + clearPendingEditMessageId(): void { + this.pendingEditMessageId = null; + } + + clearPendingMessage(convId: string): void { + this.pendingMessages.delete(convId); + } + + /** Reset per-view state when (re)mounting the empty chat screen. */ + clearUIState(): void { + this.currentResponse = ''; + } + + consumePendingDraft(): { message: string; files: ChatUploadedFile[] } | null { + if (!this.pendingDraftMessage && this.pendingDraftFiles.length === 0) return null; + + const d = { files: [...this.pendingDraftFiles], message: this.pendingDraftMessage }; + + this.pendingDraftMessage = ''; + this.pendingDraftFiles = []; + + return d; + } + + consumePendingMessage( + convId: string + ): { content: string; extras?: DatabaseMessageExtra[] } | null { + const msg = this.pendingMessages.get(convId); + + if (!msg) return null; + + this.pendingMessages.delete(convId); + + return msg; + } + + async continueAssistantMessage(messageId: string): Promise { + return this.flows.continueAssistantMessage(messageId); + } + + async createAssistantMessage(parentId?: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) throw new Error('No active conversation'); + + return await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId: activeConv.id, + model: null, + role: MessageRole.ASSISTANT, + timestamp: Date.now(), + toolCalls: '', + type: MessageType.TEXT + }, + parentId || null + ); + } + + async deleteMessage(messageId: string): Promise { + return this.flows.deleteMessage(messageId); + } + + /** + * Server-side stream sessions (discovery, attach/replay, resume retry, + * remote-running snapshot) live in ChatStreamManager. + */ + async discoverActiveStream(convId: string): Promise { + return this.streams.discoverActiveStream(convId); + } + + dismissErrorDialog(): void { + this.errorDialogState = null; + } + + async editAssistantMessage( + messageId: string, + newContent: string, + shouldBranch: boolean + ): Promise { + return this.flows.editAssistantMessage(messageId, newContent, shouldBranch); + } + + async editMessageWithBranching( + messageId: string, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ): Promise { + return this.flows.editMessageWithBranching(messageId, newContent, newExtras); + } + + async editUserMessagePreserveResponses( + messageId: string, + newContent: string, + newExtras?: DatabaseMessageExtra[] + ): Promise { + return this.flows.editUserMessagePreserveResponses(messageId, newContent, newExtras); + } + + getAddFilesHandler(): ((files: File[]) => void) | null { + return this.addFilesHandler; + } + + /** Convs with any activity (local pipe or remote session), sidebar spinners. */ + getAllLoadingChats(): string[] { + return this.activity.loadingConvs; + } + + getApiOptions(): Record { + const currentConfig = settingsStore.config; + const hasValue = (value: unknown): boolean => + value !== undefined && value !== null && value !== ''; + const apiOptions: Record = { stream: true, timings_per_token: true }; + + if (serverStore.isRouterMode) { + const modelName = modelsStore.selectedModelName; + + if (modelName) apiOptions.model = modelName; + } + + if (currentConfig.systemMessage) apiOptions.systemMessage = currentConfig.systemMessage; + + if (currentConfig.disableReasoningParsing) apiOptions.disableReasoningParsing = true; + + if (currentConfig.excludeReasoningFromContext) apiOptions.excludeReasoningFromContext = true; + + // an explicit reasoning choice overrides the server default, DEFAULT sends nothing + const effort = conversationsStore.preferences.getReasoningEffort(); + + if (effort !== ReasoningEffort.DEFAULT) { + apiOptions.enableThinking = effort !== ReasoningEffort.OFF; + + if (effort !== ReasoningEffort.OFF) apiOptions.reasoningEffort = effort; + } + + if (hasValue(currentConfig.temperature)) + apiOptions.temperature = Number(currentConfig.temperature); + + if (hasValue(currentConfig.max_tokens)) + apiOptions.max_tokens = Number(currentConfig.max_tokens); + + if (hasValue(currentConfig.dynatemp_range)) + apiOptions.dynatemp_range = Number(currentConfig.dynatemp_range); + + if (hasValue(currentConfig.dynatemp_exponent)) + apiOptions.dynatemp_exponent = Number(currentConfig.dynatemp_exponent); + + if (hasValue(currentConfig.top_k)) apiOptions.top_k = Number(currentConfig.top_k); + + if (hasValue(currentConfig.top_p)) apiOptions.top_p = Number(currentConfig.top_p); + + if (hasValue(currentConfig.min_p)) apiOptions.min_p = Number(currentConfig.min_p); + + if (hasValue(currentConfig.xtc_probability)) + apiOptions.xtc_probability = Number(currentConfig.xtc_probability); + + if (hasValue(currentConfig.xtc_threshold)) + apiOptions.xtc_threshold = Number(currentConfig.xtc_threshold); + + if (hasValue(currentConfig.typ_p)) apiOptions.typ_p = Number(currentConfig.typ_p); + + if (hasValue(currentConfig.repeat_last_n)) + apiOptions.repeat_last_n = Number(currentConfig.repeat_last_n); + + if (hasValue(currentConfig.repeat_penalty)) + apiOptions.repeat_penalty = Number(currentConfig.repeat_penalty); + + if (hasValue(currentConfig.presence_penalty)) + apiOptions.presence_penalty = Number(currentConfig.presence_penalty); + + if (hasValue(currentConfig.frequency_penalty)) + apiOptions.frequency_penalty = Number(currentConfig.frequency_penalty); + + if (hasValue(currentConfig.dry_multiplier)) + apiOptions.dry_multiplier = Number(currentConfig.dry_multiplier); + + if (hasValue(currentConfig.dry_base)) apiOptions.dry_base = Number(currentConfig.dry_base); + + if (hasValue(currentConfig.dry_allowed_length)) + apiOptions.dry_allowed_length = Number(currentConfig.dry_allowed_length); + + if (hasValue(currentConfig.dry_penalty_last_n)) + apiOptions.dry_penalty_last_n = Number(currentConfig.dry_penalty_last_n); + + if (currentConfig.samplers) apiOptions.samplers = currentConfig.samplers; + + if (hasValue(currentConfig.backend_sampling)) + apiOptions.backend_sampling = currentConfig.backend_sampling; + + if (currentConfig.customJson) apiOptions.custom = currentConfig.customJson; + + return apiOptions; + } + + getChatStreaming(convId: string): { response: string; messageId: string } | undefined { + return this.getChatStreamingState(convId); + } + + async getDeletionInfo(messageId: string): Promise<{ + totalCount: number; + userMessages: number; + assistantMessages: number; + messageTypes: string[]; + }> { + return this.flows.getDeletionInfo(messageId); + } + + getOrCreateAbortController(convId: string): AbortController { + let c = this.abortControllers.get(convId); + + if (!c || c.signal.aborted) { + c = new AbortController(); + this.abortControllers.set(convId, c); + } + + return c; + } + + getPendingMessageContent(convId: string): string | null { + return this.pendingMessages.get(convId)?.content ?? null; + } + + getPendingMessageExtras(convId: string): DatabaseMessageExtra[] | undefined { + return this.pendingMessages.get(convId)?.extras; + } + + getResumeModel(convId: string): string | null { + return this.streams.getResumeModel(convId); + } + + hasPendingDraft(): boolean { + return Boolean(this.pendingDraftMessage) || this.pendingDraftFiles.length > 0; + } + + hasPendingMessage(convId: string): boolean { + return this.pendingMessages.has(convId); + } + + injectPendingMessage(convId: string, content: string, extras?: DatabaseMessageExtra[]): void { + this.pendingMessages.set(convId, { content, extras }); + } + + isChatLoading(convId: string): boolean { + return this.activity.isLocal(convId); + } + + isChatLoadingInternal(convId: string): boolean { + return this.activity.isLocal(convId) || this.chatStreamingStates.has(convId); + } + + isEditing(): boolean { + return this.isEditModeActive; + } + + /** True while the active conversation has a live streaming pipe. */ + isStreaming(): boolean { + return this.chatStreamingStates.has(conversationsStore.activeConversation?.id ?? ''); + } + + /** + * Record a working-directory change into chat history as a synthetic + * user message, so the model sees it on its next turn (the client + * sends the cwd itself via the x-tool-cwd header on tool calls). + * A plain user message is used because some chat templates reject + * tool messages without a preceding tool call. + */ + async recordCwdChange(cwd: string | null): Promise { + const content = cwd + ? formatCwdMessage(cwd, await toolsStore.resolveServerHome()) + : CWD_CLEARED_TEXT; + // Reuse the trailing cwd row when it is already the last message, so + // repeated picks update it in place instead of stacking another row. + const last = conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1]; + + if (last && last.role === MessageRole.USER && last.isSynthetic === true) { + await DatabaseService.updateMessage(last.id, { content, isSynthetic: true }); + conversationsStore.updateMessageAtIndex(conversationsStore.activeMessages.length - 1, { + content, + isSynthetic: true + }); + + return; + } + + await this.addMessage(MessageRole.USER, content, MessageType.TEXT, '-1', undefined, true); + } + + async regenerateMessage(messageId: string): Promise { + return this.flows.regenerateMessage(messageId); + } + + async regenerateMessageWithBranching(messageId: string, modelOverride?: string): Promise { + return this.flows.regenerateMessageWithBranching(messageId, modelOverride); + } + + async removeSystemPromptPlaceholder(messageId: string): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return false; + + try { + const allMessages = await conversationsStore.getConversationMessages(activeConv.id); + const systemMessage = findMessageById(allMessages, messageId); + + if (!systemMessage || systemMessage.role !== MessageRole.SYSTEM) return false; + + const rootMessage = allMessages.find((m) => m.type === 'root' && m.parent === null); + + if (!rootMessage) return false; + + if (allMessages.length === 2 && systemMessage.children.length === 0) { + await conversationsStore.deleteConversation(activeConv.id); + + return true; + } + + for (const childId of systemMessage.children) { + await DatabaseService.updateMessage(childId, { parent: rootMessage.id }); + const childIndex = conversationsStore.findMessageIndex(childId); + + if (childIndex !== -1) + conversationsStore.updateMessageAtIndex(childIndex, { parent: rootMessage.id }); + } + await DatabaseService.updateMessage(rootMessage.id, { + children: [ + ...rootMessage.children.filter((id: string) => id !== messageId), + ...systemMessage.children + ] + }); + await DatabaseService.deleteMessage(messageId); + const systemIndex = conversationsStore.findMessageIndex(messageId); + + if (systemIndex !== -1) conversationsStore.activeMessages.splice(systemIndex, 1); + + conversationsStore.updateConversationTimestamp(); + + return false; + } catch (error) { + console.error('Failed to remove system prompt placeholder:', error); + + return false; + } + } + + savePendingDraft(message: string, files: ChatUploadedFile[]): void { + this.pendingDraftMessage = message; + this.pendingDraftFiles = [...files]; + } + async sendMessage(content: string, extras?: DatabaseMessageExtra[]): Promise { + if (!content.trim() && (!extras || extras.length === 0)) return; + + const activeConv = conversationsStore.activeConversation; + + // If agentic loop is running, inject as a steering message instead of starting a new flow + if (activeConv && agenticStore.isRunning(activeConv.id)) { + agenticStore.injectSteeringMessage(activeConv.id, content, extras); + + return; + } + + // If non-agentic streaming is active, queue as a pending message to send after completion + if (activeConv && this.isChatLoadingInternal(activeConv.id)) { + this.injectPendingMessage(activeConv.id, content, extras); + + return; + } + + // Cancel any in-flight pre-encode request + this.cancelPreEncode(); + + // Consume MCP resource attachments - converts them to extras and clears the live store + const resourceExtras = mcpStore.consumeResourceAttachmentsAsExtras(); + const allExtras = resourceExtras.length > 0 ? [...(extras || []), ...resourceExtras] : extras; + + let isNewConversation = false; + + if (!activeConv) { + await conversationsStore.createConversation(); + isNewConversation = true; + } + + const currentConv = conversationsStore.activeConversation; + + if (!currentConv) return; + + this.showErrorDialog(null); + this.setChatLoading(currentConv.id, true); + this.clearChatStreaming(currentConv.id); + try { + let parentIdForUserMessage: string | undefined; + + if (isNewConversation) { + const rootId = await DatabaseService.createRootMessage(currentConv.id); + const currentConfig = settingsStore.config; + const systemPrompt = currentConfig.systemMessage?.toString().trim(); + + let sysOrRootId = rootId; + + if (systemPrompt) { + const systemMessage = await DatabaseService.createSystemMessage( + currentConv.id, + systemPrompt, + rootId + ); + + conversationsStore.addMessageToActive(systemMessage); + sysOrRootId = systemMessage.id; + } + + // Reflect a working directory picked on the new-chat screen into + // chat history before the first user message, so the model sees + // it on its first turn. createConversation() has already threaded + // the pending pick onto the conversation. + if (currentConv.cwd) { + const cwdMessage = await this.addMessage( + MessageRole.USER, + formatCwdMessage(currentConv.cwd, await toolsStore.resolveServerHome()), + MessageType.TEXT, + sysOrRootId, + undefined, + true + ); + + parentIdForUserMessage = cwdMessage.id; + } else { + parentIdForUserMessage = sysOrRootId; + } + } + + const userMessage = await this.addMessage( + MessageRole.USER, + content, + MessageType.TEXT, + parentIdForUserMessage ?? '-1', + allExtras + ); + + if (isNewConversation && content) + await conversationsStore.applyTitleFromContent(currentConv.id, content); + + const assistantMessage = await this.createAssistantMessage(userMessage.id); + + conversationsStore.addMessageToActive(assistantMessage); + await this.streamChatCompletion( + conversationsStore.activeMessages.slice(0, -1), + assistantMessage, + undefined, + undefined, + undefined, + settingsStore.config.titleGenerationUseLLM && isNewConversation ? content : undefined + ); + } catch (error) { + if (isAbortError(error)) { + this.setChatLoading(currentConv.id, false); + + return; + } + + console.error('Failed to send message:', error); + this.setChatLoading(currentConv.id, false); + const dialogType = + error instanceof Error && error.name === 'TimeoutError' + ? ErrorDialogType.TIMEOUT + : ErrorDialogType.SERVER; + const contextInfo = ( + error as Error & { contextInfo?: { n_prompt_tokens: number; n_ctx: number } } + ).contextInfo; + + this.showErrorDialog({ + contextInfo, + message: error instanceof Error ? error.message : 'Unknown error', + type: dialogType + }); + } + } + + setChatLoading(convId: string, loading: boolean): void { + if (loading) { + this.activity.markLocal(convId); + } else { + this.activity.localEnded(convId); + this.setChatReasoning(convId, false); + } + } + + setChatReasoning(convId: string, reasoning: boolean): void { + if (reasoning) this.chatReasoningStates.set(convId, true); + else this.chatReasoningStates.delete(convId); + } + + setChatStreaming( + convId: string, + response: string, + messageId: string, + model?: string | null + ): void { + this.chatStreamingStates.set(convId, { + messageId, + model: model ?? this.chatStreamingStates.get(convId)?.model, + response + }); + + if (convId === conversationsStore.activeConversation?.id) this.currentResponse = response; + } + + setEditModeActive(handler: (files: File[]) => void): void { + this.isEditModeActive = true; + this.addFilesHandler = handler; + } + + showErrorDialog(state: ErrorDialogState | null): void { + this.errorDialogState = state; + } + + async stopGeneration(): Promise { + const activeConv = conversationsStore.activeConversation; + + if (!activeConv) return; + + await this.stopGenerationForChat(activeConv.id); + } + + async stopGenerationForChat(convId: string): Promise { + await this.savePartialResponseIfNeeded(convId); + // tell the server to stop the generation, not just drop the HTTP socket. without this the + // detached drain keeps producing tokens until eos or max_tokens. use the frozen identity + // captured when the session started, not the live dropdown + const streamStateForStop = this.chatStreamingStates.get(convId); + const modelForStop = streamStateForStop?.model ?? ChatService.getStreamState(convId)?.model; + + void ChatService.cancelServerStream(convId, modelForStop); + // an explicit stop leaves nothing to resume and kills a pending resume retry + ChatService.clearStreamState(convId); + this.streams.cancelResumeRetry(convId); + this.abortRequest(convId); + this.setChatLoading(convId, false); + this.clearChatStreaming(convId); + this.processing.setState(convId, null); + this.clearPendingMessage(convId); + } + + async streamChatCompletion( + allMessages: DatabaseMessage[], + assistantMessage: DatabaseMessage, + onComplete?: (content: string) => Promise, + onError?: (error: Error) => void, + modelOverride?: string | null, + firstUserMessageContent?: string + ): Promise { + // the ::model suffix in the stream identity is only for router mode, where it routes to the + // owning child. in single-model mode the identity stays the bare conv id so that attach, stop + // and reattach all agree, regardless of fresh send vs regenerate passing a resolved model + let effectiveModel: string | null | undefined = undefined; + + if (serverStore.isRouterMode) { + const conversationModel = getConversationModel(allMessages); + + effectiveModel = modelOverride || modelsStore.selectedModelName || conversationModel; + } + + if (serverStore.isRouterMode && effectiveModel) { + if (!modelsStore.props.getModelProps(effectiveModel)) + await modelsStore.props.fetchModelProps(effectiveModel); + } + + // Mutable state for the current message being streamed + let currentMessageId = assistantMessage.id; + let streamedContent = ''; + let streamedReasoningContent = ''; + let resolvedModel: string | null = null; + let modelPersisted = false; + + const convId = assistantMessage.convId; + + // Tracks the last message created in this flow. Used as the parent for the next + // turn's assistant message so createAssistantMessage does not have to read + // conversationsStore.activeMessages, which may belong to a different conversation + // after the user navigates while the loop is still running. + let lastCreatedInFlow = currentMessageId; + + // freeze the POST identity from t0 so a stop cancels with the exact session key, + // never a stale or empty model resolved later + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); + + const recordModel = (modelName: string | null | undefined, persistImmediately = true): void => { + if (!modelName) return; + + const n = normalizeModelName(modelName); + + if (!n || n === resolvedModel) return; + + resolvedModel = n; + const idx = conversationsStore.findMessageIndex(currentMessageId); + + conversationsStore.updateMessageAtIndex(idx, { model: n }); + + if (persistImmediately && !modelPersisted) { + modelPersisted = true; + DatabaseService.updateMessage(currentMessageId, { model: n }).catch(() => { + modelPersisted = false; + resolvedModel = null; + }); + } + }; + + let completionIdRecorded = false; + + const recordCompletionId = (id: string): void => { + if (!id || completionIdRecorded) return; + + completionIdRecorded = true; + const idx = conversationsStore.findMessageIndex(currentMessageId); + + conversationsStore.updateMessageAtIndex(idx, { completionId: id }); + DatabaseService.updateMessage(currentMessageId, { completionId: id }).catch(() => { + completionIdRecorded = false; + }); + }; + const updateStreamingUI = () => { + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); + const idx = conversationsStore.findMessageIndex(currentMessageId); + + conversationsStore.updateMessageAtIndex(idx, { content: streamedContent }); + }; + const cleanupStreamingState = () => { + this.setChatLoading(convId, false); + this.clearChatStreaming(convId, currentMessageId); + this.processing.setState(convId, null); + }; + + this.processing.setActiveConversation(convId); + const abortController = this.getOrCreateAbortController(convId); + const streamCallbacks: ChatStreamCallbacks = { + createAssistantMessage: async () => { + // Reset streaming state for new message + streamedContent = ''; + streamedReasoningContent = ''; + + const msg = await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId, + model: resolvedModel, + role: MessageRole.ASSISTANT, + timestamp: Date.now(), + toolCalls: '', + type: MessageType.TEXT + }, + lastCreatedInFlow + ); + + if (conversationsStore.activeConversation?.id === convId) { + conversationsStore.addMessageToActive(msg); + } + + currentMessageId = msg.id; + lastCreatedInFlow = msg.id; + + return msg; + }, + createToolResultMessage: async ( + toolCallId: string, + content: string, + extras?: DatabaseMessageExtra[], + toolCwd?: string + ) => { + const msg = await DatabaseService.createMessageBranch( + { + children: [], + content, + convId, + extra: extras, + role: MessageRole.TOOL, + timestamp: Date.now(), + toolCallId, + toolCalls: '', + toolCwd, + type: MessageType.TEXT + }, + currentMessageId + ); + + // mirror into the active store and move the node pointer only when this + // conversation is displayed; otherwise persist the node move straight to + // the db for the owning conv so a foreign conv's currNode stays untouched + if (conversationsStore.activeConversation?.id === convId) { + conversationsStore.addMessageToActive(msg); + await conversationsStore.updateCurrentNode(msg.id); + } else { + await DatabaseService.updateCurrentNode(convId, msg.id); + } + + lastCreatedInFlow = msg.id; + + return msg; + }, + onAssistantTurnComplete: async ( + content: string, + reasoningContent: string | undefined, + timings: ChatMessageTimings | undefined, + toolCalls: import('$lib/types/api').ApiChatCompletionToolCall[] | undefined + ) => { + const updateData: Record = { + content, + reasoningContent: reasoningContent || undefined, + timings, + toolCalls: toolCalls ? JSON.stringify(toolCalls) : '' + }; + + if (resolvedModel && !modelPersisted) updateData.model = resolvedModel; + + await DatabaseService.updateMessage(currentMessageId, updateData); + const idx = conversationsStore.findMessageIndex(currentMessageId); + const uiUpdate: Partial = { + content, + reasoningContent: reasoningContent || undefined, + toolCalls: toolCalls ? JSON.stringify(toolCalls) : '' + }; + + if (timings) uiUpdate.timings = timings; + + if (resolvedModel) uiUpdate.model = resolvedModel; + + // touch the active ui array and node pointer only when this conversation + // is displayed; otherwise persist the node move straight to the db so a + // foreign conv's currNode stays untouched + if (conversationsStore.activeConversation?.id === convId) { + conversationsStore.updateMessageAtIndex(idx, uiUpdate); + await conversationsStore.updateCurrentNode(currentMessageId); + } else { + await DatabaseService.updateCurrentNode(convId, currentMessageId); + } + }, + onAttachments: (messageId: string, extras: DatabaseMessageExtra[]) => { + if (!extras.length) return; + + const idx = conversationsStore.findMessageIndex(messageId); + + if (idx === -1) return; + + const msg = conversationsStore.activeMessages[idx]; + const updatedExtras = [...(msg.extra || []), ...extras]; + + conversationsStore.updateMessageAtIndex(idx, { extra: updatedExtras }); + DatabaseService.updateMessage(messageId, { extra: updatedExtras }).catch(console.error); + }, + onChunk: (chunk: string) => { + streamedContent += chunk; + updateStreamingUI(); + this.setChatReasoning(convId, false); + }, + onCompletionId: (id: string) => recordCompletionId(id), + onError: async (error: Error) => { + if (isAbortError(error)) { + cleanupStreamingState(); + // If aborted with a pending message (e.g. "Send immediately"), re-send it + const pending = this.consumePendingMessage(convId); + + if (pending) { + this.sendMessage(pending.content, pending.extras); + } + + return; + } + + console.error('Streaming error:', error); + // keep whatever was streamed so far, the message stays in memory and in DB + await this.savePartialResponseIfNeeded(convId); + cleanupStreamingState(); + this.clearPendingMessage(convId); + + const contextInfo = ( + error as Error & { contextInfo?: { n_prompt_tokens: number; n_ctx: number } } + ).contextInfo; + + this.showErrorDialog({ + contextInfo, + message: error.message, + type: error.name === 'TimeoutError' ? ErrorDialogType.TIMEOUT : ErrorDialogType.SERVER + }); + + if (onError) onError(error); + }, + onFlowComplete: (finalTimings?: ChatMessageTimings) => { + if (finalTimings) { + const idx = conversationsStore.findMessageIndex(assistantMessage.id); + + conversationsStore.updateMessageAtIndex(idx, { timings: finalTimings }); + DatabaseService.updateMessage(assistantMessage.id, { + timings: finalTimings + }).catch(console.error); + } + + cleanupStreamingState(); + + if (onComplete) onComplete(streamedContent); + + if (serverStore.isRouterMode) modelsStore.fetchRouterModels().catch(console.error); + + // Pre-encode conversation in KV cache for faster next turn + if (settingsStore.config.preEncodeConversation) { + this.triggerPreEncode( + allMessages, + assistantMessage, + streamedContent, + effectiveModel, + !!settingsStore.config.excludeReasoningFromContext + ); + } + }, + onModel: (modelName: string) => recordModel(modelName), + onReasoningChunk: (chunk: string) => { + streamedReasoningContent += chunk; + // mark streaming state so a stop mid-thinking can persist the partial reasoning + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); + const idx = conversationsStore.findMessageIndex(currentMessageId); + + conversationsStore.updateMessageAtIndex(idx, { + reasoningContent: streamedReasoningContent + }); + this.setChatReasoning(convId, true); + }, + onTimings: (timings?: ChatMessageTimings, promptProgress?: ChatMessagePromptProgress) => { + this.processing.applyStreamTimings(timings, promptProgress, convId); + }, + onToolCallsStreaming: (toolCalls) => { + const idx = conversationsStore.findMessageIndex(currentMessageId); + + conversationsStore.updateMessageAtIndex(idx, { + toolCalls: JSON.stringify(toolCalls) + }); + }, + onTurnComplete: (intermediateTimings: ChatMessageTimings) => { + // Update the first assistant message with cumulative agentic timings + const idx = conversationsStore.findMessageIndex(assistantMessage.id); + + conversationsStore.updateMessageAtIndex(idx, { timings: intermediateTimings }); + }, + updateToolResultMessage: async ( + messageId: string, + content: string, + extras?: DatabaseMessageExtra[] + ) => { + // Persist latest content + merged extras; mirror into the active + // store so the chat view sees live updates for streaming tools + // (e.g. exec_shell_command). The existing tool message node + // pointer stays put - the renderer is already scoped to it. + const updates: Partial = { content }; + + if (extras) { + const idx = conversationsStore.findMessageIndex(messageId); + const existing = idx >= 0 ? (conversationsStore.activeMessages[idx]?.extra ?? []) : []; + const merged = [...existing, ...extras]; + + updates.extra = merged; + } + + if (conversationsStore.activeConversation?.id === convId) { + const idx = conversationsStore.findMessageIndex(messageId); + + if (idx >= 0) conversationsStore.updateMessageAtIndex(idx, updates); + } + + await DatabaseService.updateMessage(messageId, updates); + } + }; + const perChatOverrides = conversationsStore.preferences.getAllMcpServerOverrides(); + + { + const agenticResult = await agenticStore.runAgenticFlow({ + callbacks: streamCallbacks, + conversationId: convId, + flowRootMessageId: assistantMessage.id, + messages: allMessages, + options: { + ...this.getApiOptions(), + ...(effectiveModel ? { model: effectiveModel } : {}) + }, + perChatOverrides, + signal: abortController.signal + }); + + if (agenticResult.handled) { + // Generate LLM based title for new conversations after agentic flow completes + if (firstUserMessageContent) { + await this.generateTitleWithLLM(firstUserMessageContent, streamedContent, convId); + } + + // Check if there's a pending steering message to re-send + const pending = agenticStore.consumePendingSteeringMessage(convId); + + if (pending) { + await this.sendMessage(pending.content, pending.extras); + } + + return; + } + } + + await ChatService.sendMessage( + allMessages, + { + ...this.getApiOptions(), + ...(effectiveModel ? { model: effectiveModel } : {}), + onChunk: streamCallbacks.onChunk, + onComplete: async ( + finalContent?: string, + reasoningContent?: string, + timings?: ChatMessageTimings, + toolCalls?: string + ) => { + const content = streamedContent || finalContent || ''; + const reasoning = streamedReasoningContent || reasoningContent; + const updateData: Record = { + content, + reasoningContent: reasoning || undefined, + timings, + toolCalls: toolCalls || '' + }; + + if (resolvedModel && !modelPersisted) updateData.model = resolvedModel; + + await DatabaseService.updateMessage(currentMessageId, updateData); + const idx = conversationsStore.findMessageIndex(currentMessageId); + const uiUpdate: Partial = { + content, + reasoningContent: reasoning || undefined, + toolCalls: toolCalls || '' + }; + + if (timings) uiUpdate.timings = timings; + + if (resolvedModel) uiUpdate.model = resolvedModel; + + conversationsStore.updateMessageAtIndex(idx, uiUpdate); + await conversationsStore.updateCurrentNode(currentMessageId); + cleanupStreamingState(); + + if (onComplete) await onComplete(content); + + if (serverStore.isRouterMode) modelsStore.fetchRouterModels().catch(console.error); + + // Generate LLM based title for new conversations (avoids stale reference + // issue when user switches conversations while streaming) + if (firstUserMessageContent) { + await this.generateTitleWithLLM(firstUserMessageContent, streamedContent, convId); + } + + // Check if there's a pending message queued during streaming + const pending = this.consumePendingMessage(convId); + + if (pending) { + await this.sendMessage(pending.content, pending.extras); + } + }, + onCompletionId: streamCallbacks.onCompletionId, + onConnectionState: (state: StreamConnectionState) => { + if (convId === conversationsStore.activeConversation?.id) { + this.streamConnectionState = state; + } + }, + onError: streamCallbacks.onError, + onModel: streamCallbacks.onModel, + onReasoningChunk: streamCallbacks.onReasoningChunk, + onTimings: streamCallbacks.onTimings, + stream: true + }, + convId, + abortController.signal + ); + } + + syncLoadingStateForChat(convId: string): void { + const s = this.chatStreamingStates.get(convId); + + this.currentResponse = s?.response || ''; + this.processing.setActiveConversation(convId); + + // Sync streaming content to activeMessages so UI displays current content + if (s?.response && s?.messageId) { + const idx = conversationsStore.findMessageIndex(s.messageId); + + if (idx !== -1) { + conversationsStore.updateMessageAtIndex(idx, { content: s.response }); + } + } + } + + async syncRemoteRunningStreams(): Promise { + return this.streams.syncRemoteRunningStreams(); + } + + /** + * Message flows (edit / regenerate / continue / delete) live in + * ChatMessageFlows; these delegate so consumers keep a single entry point. + */ + async updateMessage(messageId: string, newContent: string): Promise { + return this.flows.updateMessage(messageId, newContent); + } + private abortRequest(convId?: string): void { + if (convId) { + const c = this.abortControllers.get(convId); + + if (c) { + c.abort(); + this.abortControllers.delete(convId); + } + } else { + for (const c of this.abortControllers.values()) c.abort(); + this.abortControllers.clear(); + } + } + + private async generateTitleWithLLM( + userContent: string, + assistantContent: string, + convId: string + ): Promise { + const effectiveModel = + serverStore.isRouterMode && modelsStore.selectedModelName + ? modelsStore.selectedModelName + : undefined; + const configValue = settingsStore.config; + const titlePromptTemplate = + typeof configValue.titleGenerationPrompt === 'string' && + configValue.titleGenerationPrompt.trim() + ? configValue.titleGenerationPrompt + : TITLE_GENERATION.DEFAULT_PROMPT; + const titlePrompt = titlePromptTemplate + .replace('{{USER}}', String(userContent || '')) + .replace('{{ASSISTANT}}', String(assistantContent || '')); + const titleMessage: ApiChatMessageData = { + content: titlePrompt, + role: MessageRole.USER + }; + const titleResponse = await ChatService.generateTitle(titleMessage, effectiveModel); + + if (!titleResponse) { + return; + } + + let cleanTitle = titleResponse.trim(); + + cleanTitle = cleanTitle + .replace(TITLE_GENERATION.PREFIX_PATTERN, '') + .replace(TITLE_GENERATION.QUOTE_PATTERN, '') + .trim(); + + if (!cleanTitle || cleanTitle.length < TITLE_GENERATION.MIN_LENGTH) { + const firstLine = userContent.split('\n').find((l) => l.trim().length > 0); + + cleanTitle = firstLine ? firstLine.trim() : TITLE_GENERATION.FALLBACK; + } + + if (cleanTitle && cleanTitle.length >= TITLE_GENERATION.MIN_LENGTH) { + await conversationsStore.updateConversationName(convId, cleanTitle); + } + } + + private getChatStreamingState( + convId: string + ): { response: string; messageId: string } | undefined { + return this.chatStreamingStates.get(convId); + } + + private async savePartialResponseIfNeeded(convId?: string): Promise { + const conversationId = convId || conversationsStore.activeConversation?.id; + + if (!conversationId) return; + + const streamingState = this.getChatStreamingState(conversationId); + + if (!streamingState) return; + + const messages = + conversationId === conversationsStore.activeConversation?.id + ? conversationsStore.activeMessages + : await conversationsStore.getConversationMessages(conversationId); + + if (!messages.length) return; + + const lastMessage = messages[messages.length - 1]; + + if (lastMessage?.role !== MessageRole.ASSISTANT) return; + + const partialContent = streamingState.response; + const partialReasoning = lastMessage.reasoningContent || ''; + // snapshot the streamed tool calls before clearing so we still know whether + // anything was captured when deciding to skip the DB write below + const hadPartialToolCalls = !!lastMessage.toolCalls?.trim(); + + // nothing to persist when content, reasoning, and streamed tool calls are all empty + // (e.g. stop before any token). otherwise drop the partial tool call and write whatever + // was streamed: incomplete arguments (truncated JSON, missing closing quote) would + // otherwise be re-sent to the server on the next turn and rejected. + if (!partialContent.trim() && !partialReasoning.trim() && !hadPartialToolCalls) return; + + try { + const updateData: { + content?: string; + reasoningContent?: string; + toolCalls?: string; + timings?: ChatMessageTimings; + } = { + toolCalls: '' + }; + + if (partialContent.trim()) updateData.content = partialContent; + + if (partialReasoning.trim()) updateData.reasoningContent = partialReasoning; + + const lastKnownState = this.processing.getState(conversationId); + + if (lastKnownState) { + updateData.timings = { + cache_n: lastKnownState.cacheTokens || 0, + predicted_ms: + lastKnownState.tokensPerSecond && lastKnownState.tokensDecoded + ? (lastKnownState.tokensDecoded / lastKnownState.tokensPerSecond) * 1000 + : undefined, + predicted_n: lastKnownState.tokensDecoded || 0, + prompt_ms: lastKnownState.promptMs, + prompt_n: lastKnownState.promptTokens || 0 + }; + } + + await DatabaseService.updateMessage(lastMessage.id, updateData); + lastMessage.content = partialContent; + // mirror the drop into the in-memory message so the next request sent via + // sendMessage (queued pending, Send immediately, or manual follow-up) reads + // the cleared value, not whatever the streaming widget had been showing + lastMessage.toolCalls = ''; + + if (updateData.timings) lastMessage.timings = updateData.timings; + } catch (error) { + lastMessage.content = partialContent; + lastMessage.toolCalls = ''; + console.error('Failed to save partial response:', error); + } + } + + private async triggerPreEncode( + allMessages: DatabaseMessage[], + assistantMessage: DatabaseMessage, + assistantContent: string, + model?: string | null, + excludeReasoning?: boolean + ): Promise { + this.cancelPreEncode(); + this.preEncodeAbortController = new AbortController(); + + const signal = this.preEncodeAbortController.signal; + + try { + const allIdle = await ChatService.areAllSlotsIdle(model, signal); + + if (!allIdle || signal.aborted) return; + + const messagesWithAssistant: DatabaseMessage[] = [ + ...allMessages, + { ...assistantMessage, content: assistantContent } + ]; + + await ChatService.preEncode(messagesWithAssistant, model, excludeReasoning, signal); + } catch (err) { + if (!isAbortError(err)) { + console.warn('[ChatStore] Pre-encode failed:', err); + } + } + } +} + +export const chatStore = new ChatStore(); diff --git a/tools/ui/src/lib/stores/chat/processing.svelte.ts b/tools/ui/src/lib/stores/chat/processing.svelte.ts new file mode 100644 index 000000000000..69c1a6925669 --- /dev/null +++ b/tools/ui/src/lib/stores/chat/processing.svelte.ts @@ -0,0 +1,188 @@ +/** + * chatProcessingStore - Per-conversation processing state + * + * Owns the live processing snapshot shown while a conversation streams: + * token counts, tokens/sec, prompt progress. Updated from stream timings, + * restored from persisted message timings when a conversation loads. + * + * Composed under chatStore.processing; not exported from the stores barrel. + */ + +import { MessageRole } from '$lib/enums'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { serverStore } from '$lib/stores/server.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; +import type { + ApiProcessingState, + ChatMessagePromptProgress, + ChatMessageTimings, + DatabaseMessage +} from '$lib/types'; +import { SvelteMap } from 'svelte/reactivity'; + +interface ProcessingTimingData { + cache_n: number; + predicted_n: number; + predicted_per_second: number; + prompt_ms?: number; + prompt_n: number; + prompt_progress?: ChatMessagePromptProgress; +} + +export class ChatProcessingStore { + private _activeConversationId = $state(null); + private states = new SvelteMap(); + + /** Processing state of the conversation currently shown in the UI. */ + activeState = $derived( + this._activeConversationId ? (this.states.get(this._activeConversationId) ?? null) : null + ); + + get activeConversationId(): string | null { + return this._activeConversationId; + } + + /** + * Applies a stream timings event (tokens/sec + token counts) to the given + * conversation's processing state. Shared by the chat and continue flows. + */ + applyStreamTimings( + timings?: ChatMessageTimings, + promptProgress?: ChatMessagePromptProgress, + conversationId?: string + ): void { + const tokensPerSecond = + timings?.predicted_ms && timings?.predicted_n + ? (timings.predicted_n / timings.predicted_ms) * 1000 + : 0; + + this.updateFromTimings( + { + cache_n: timings?.cache_n || 0, + predicted_n: timings?.predicted_n || 0, + predicted_per_second: tokensPerSecond, + prompt_ms: timings?.prompt_ms, + prompt_n: timings?.prompt_n || 0, + prompt_progress: promptProgress + }, + conversationId + ); + } + + getConversationIds(): string[] { + return Array.from(this.states.keys()); + } + + getState(conversationId: string): ApiProcessingState | null { + return this.states.get(conversationId) ?? null; + } + + restoreFromMessages(messages: DatabaseMessage[], conversationId: string): void { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + + if (message.role === MessageRole.ASSISTANT && message.timings) { + this.setState( + conversationId, + this.parseTimingData({ + cache_n: message.timings.cache_n || 0, + predicted_n: message.timings.predicted_n || 0, + predicted_per_second: + message.timings.predicted_n && message.timings.predicted_ms + ? (message.timings.predicted_n / message.timings.predicted_ms) * 1000 + : 0, + prompt_ms: message.timings.prompt_ms, + prompt_n: message.timings.prompt_n || 0 + }) + ); + + return; + } + } + } + + setActiveConversation(conversationId: string | null): void { + this._activeConversationId = conversationId; + } + + /** Passing null clears the state for the conversation. */ + setState(conversationId: string, state: ApiProcessingState | null): void { + if (state === null) this.states.delete(conversationId); + else this.states.set(conversationId, state); + } + + updateFromTimings(timingData: ProcessingTimingData, conversationId?: string): void { + const targetId = conversationId || this._activeConversationId; + + if (targetId) { + this.setState(targetId, this.parseTimingData(timingData)); + } + } + + private getContextTotal(): number | null { + const activeConvId = this._activeConversationId; + const activeState = activeConvId ? this.getState(activeConvId) : null; + + if (activeState && typeof activeState.contextTotal === 'number' && activeState.contextTotal > 0) + return activeState.contextTotal; + + if (serverStore.isRouterMode) { + const modelContextSize = modelsStore.selectedModelContextSize; + + if (typeof modelContextSize === 'number' && modelContextSize > 0) { + return modelContextSize; + } + } else { + const propsContextSize = serverStore.contextSize; + + if (typeof propsContextSize === 'number' && propsContextSize > 0) { + return propsContextSize; + } + } + + return null; + } + + private parseTimingData(timingData: ProcessingTimingData): ApiProcessingState { + const cacheTokens = timingData.cache_n || 0, + predictedTokens = timingData.predicted_n || 0, + promptMs = timingData.prompt_ms || undefined, + promptTokens = timingData.prompt_n || 0, + tokensPerSecond = timingData.predicted_per_second || 0; + const promptProgress = timingData.prompt_progress; + const contextTotal = this.getContextTotal(); + const currentConfig = settingsStore.config; + const outputTokensMax = currentConfig.max_tokens || -1; + const contextUsed = promptTokens + cacheTokens + predictedTokens, + outputTokensUsed = predictedTokens; + const progressCache = promptProgress?.cache || 0, + progressActualDone = (promptProgress?.processed ?? 0) - progressCache, + progressActualTotal = (promptProgress?.total ?? 0) - progressCache; + const progressPercent = promptProgress + ? Math.round((progressActualDone / progressActualTotal) * 100) + : undefined; + + return { + cacheTokens, + contextTotal, + contextUsed, + hasNextToken: predictedTokens > 0, + outputTokensMax, + outputTokensUsed, + progressPercent, + promptMs, + promptProgress, + promptTokens, + speculative: false, + status: predictedTokens > 0 ? 'generating' : promptProgress ? 'preparing' : 'idle', + temperature: currentConfig.temperature ?? 0.8, + tokensDecoded: predictedTokens, + tokensPerSecond, + tokensRemaining: outputTokensMax - predictedTokens, + topP: currentConfig.top_p ?? 0.95 + }; + } +} + +export const chatProcessingStore = new ChatProcessingStore(); diff --git a/tools/ui/src/lib/stores/chat/streams.svelte.ts b/tools/ui/src/lib/stores/chat/streams.svelte.ts new file mode 100644 index 000000000000..5abbc81fb631 --- /dev/null +++ b/tools/ui/src/lib/stores/chat/streams.svelte.ts @@ -0,0 +1,494 @@ +/** + * ChatStreamManager - Server-side stream sessions for conversations + * + * Owns the attach lifecycle for streams that live on the server: discovery, + * replay from byte 0, and resume retry while the owning model loads. The + * remote-running snapshot it produces feeds the chat activity ledger + * (chatStore.activity), which owns the actual running-conv state. Created + * and owned by chatStore; the host exposes the per-conversation state setters. + */ + +import { CONVERSATION_ID_SEPARATOR, STREAM_RESUME_RETRY_MS } from '$lib/constants'; +import { MessageRole, MessageType, StreamConnectionState } from '$lib/enums'; +import { ChatService } from '$lib/services/chat.service'; +import { DatabaseService } from '$lib/services/database.service'; +import type { ChatActivityStore } from '$lib/stores/chat/activity.svelte'; +import type { ChatProcessingStore } from '$lib/stores/chat/processing.svelte'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import type { ApiStreamSession, ChatMessageTimings, DatabaseMessage } from '$lib/types'; +import { streamIdentity } from '$lib/utils'; +import { SvelteMap, SvelteSet } from 'svelte/reactivity'; + +/** + * The slice of chatStore the manager drives. Kept narrow on purpose so the + * manager cannot reach around the host's full surface; chatStore implements + * this structurally. + */ +export interface ChatStreamHost { + activity: ChatActivityStore; + processing: ChatProcessingStore; + chatStreamingStates: SvelteMap< + string, + { response: string; messageId: string; model?: string | null } + >; + streamConnectionState: StreamConnectionState; + getOrCreateAbortController(convId: string): AbortController; + setChatLoading(convId: string, loading: boolean): void; + setChatStreaming( + convId: string, + response: string, + messageId: string, + model?: string | null + ): void; + clearChatStreaming(convId: string, messageId?: string): void; +} + +export class ChatStreamManager { + // in-flight discoverActiveStream guard, keyed by conv id + private discoveringConvs = new SvelteSet(); + // convs whose resume waits on a model load: their loading state belongs to the retry loop, + // so discoverActiveStream must not treat it as a live send and bail + private resumePendingConvs = new SvelteSet(); + // pending resume retry timers while an owning model loads, one per conv + private resumeRetryTimers = new SvelteMap>(); + + /** Kill a pending resume retry, e.g. on explicit stop. */ + cancelResumeRetry(convId: string): void { + const timer = this.resumeRetryTimers.get(convId); + + if (timer !== undefined) { + clearTimeout(timer); + this.resumeRetryTimers.delete(convId); + } + + this.resumePendingConvs.delete(convId); + } + + constructor(private host: ChatStreamHost) {} + + async discoverActiveStream(convId: string): Promise { + if (!convId) return; + + if (this.host.chatStreamingStates.has(convId)) return; + + if (this.host.activity.isLocal(convId) && !this.resumePendingConvs.has(convId)) return; + + // concurrency guard: another discover may already be running for this conv (typical race + // between mount and visibilitychange on tab switch). a second concurrent fetch on the same + // /v1/stream would duplicate every byte into the DB message, this guard bounces it + if (this.discoveringConvs.has(convId)) return; + + this.discoveringConvs.add(convId); + + try { + // the model is frozen at POST time, rebuild the exact conv::model identity from the + // persisted state so the lookup key matches what the server stored. null means a single + // model conv with no ::suffix, only guess from the dropdown with no persisted state + const localState = ChatService.getStreamState(convId); + const streamId = ChatService.resumeStreamIdentity( + convId, + localState, + modelsStore.selectedModelName + ); + // primary path: ask the server which sessions exist for this identity + const serverTarget = await this.probeServerStream(streamId); + + if (serverTarget) { + // pass the full server side identity (may carry a ::model suffix) so the GET routes + // straight to the owning session, no probe or fan out + await this.attachServerStream(convId, serverTarget.conversation_id); + + return; + } + + // fallback: local state remembers an interrupted byte offset for this conv, the server may + // still have a live session matching that identity (we just lost the bytes mid stream). retry + // with the frozen identity, the server probe inside attachServerStream tells us if it exists + if (!localState) { + return; + } + + // quiet status probe first: a full attach flips the loading UI on every try, probing + // keeps the retry loop invisible while the owning model is still loading (503) + const status = await ChatService.probeResumeStatus(streamId); + + if (status === 503) { + // make the wait visible: the empty assistant row persisted at send time renders + // the processing info, whose model load percentage flows from the models feed + this.resumePendingConvs.add(convId); + this.host.setChatLoading(convId, true); + + if (!this.resumeRetryTimers.has(convId)) { + this.resumeRetryTimers.set( + convId, + setTimeout(() => { + this.resumeRetryTimers.delete(convId); + void this.discoverActiveStream(convId); + }, STREAM_RESUME_RETRY_MS) + ); + } + + return; + } + + if (this.resumePendingConvs.delete(convId) && status !== 200) { + // the wait is over without a session to attach, drop the visible loading state + this.host.setChatLoading(convId, false); + } + + if (status === 0) { + // transient network failure, the next mount or visibility change retries + return; + } + + if (status !== 200) { + // the session is gone (stopped, TTL expired), nothing to resume anymore + ChatService.clearStreamState(convId); + + return; + } + + await this.attachServerStream(convId, streamId); + + // if attachServerStream failed (session gone, TTL expired), clear the local state to avoid retrying forever + if (!this.host.chatStreamingStates.has(convId) && !this.host.activity.isLocal(convId)) { + ChatService.clearStreamState(convId); + } + } finally { + this.discoveringConvs.delete(convId); + } + } + + /** + * Model frozen at send time for a stream awaiting resume, from the persisted stream state. + * The load progress indicator targets it after a reload, when the message row has no model + * yet and the dropdown selection may not be restored. + */ + getResumeModel(convId: string): string | null { + return ChatService.getStreamState(convId)?.model ?? null; + } + + /** + * Resync the activity ledger's remote set from the backend. Called by the layout at mount and + * on visibilitychange, no polling. A snapshot semantic: stale entries for sessions that + * finalized while the browser was elsewhere are dropped naturally. + */ + async syncRemoteRunningStreams(): Promise { + // the conversations store loads from IndexedDB asynchronously, the +layout onMount caller + // fires before that finishes. read ids straight from the DB so the result does not depend + // on the store init race, and the sidebar spinners light up at first paint for every conv + // the user owns even if it has not been hydrated into the store yet + let ids: string[]; + + try { + const all = await DatabaseService.getAllConversations(); + + ids = all.map((c) => c.id).filter((id) => !!id); + } catch (e) { + console.warn('syncRemoteRunningStreams DB read failed:', e); + + return; + } + + // only ask about conv ids the user already owns + if (ids.length === 0) { + this.host.activity.applyRemoteSnapshot([]); + + return; + } + + // rebuild the frozen conv::model identity per conv so a session started with a model still + // matches. the server response is mapped back to the bare id below for the sidebar set + const lookupIds = ids.map((id) => + ChatService.resumeStreamIdentity(id, ChatService.getStreamState(id), null) + ); + + let sessions: ApiStreamSession[]; + + try { + sessions = await ChatService.lookupStreamSessions(lookupIds); + } catch (e) { + console.warn('syncRemoteRunningStreams lookup failed:', e); + + return; + } + const running = new SvelteSet(); + + for (const s of sessions) { + if (s && !s.is_done && typeof s.conversation_id === 'string' && s.conversation_id) { + // strip the optional ::model suffix, the sidebar set is keyed by the bare conv id + const sepIdx = s.conversation_id.indexOf(CONVERSATION_ID_SEPARATOR); + const bareId = sepIdx === -1 ? s.conversation_id : s.conversation_id.slice(0, sepIdx); + + running.add(bareId); + } + } + this.host.activity.applyRemoteSnapshot(running); + } + + private async attachServerStream(convId: string, streamId?: string): Promise { + if (!convId) return; + + if (this.host.chatStreamingStates.has(convId)) return; + + // flip the spinner immediately, the user sees activity as soon as the conv becomes active + this.host.setChatLoading(convId, true); + + // only set the active processing conv if we are looking at it, otherwise a background + // attach would steal the indicator from the conv the user is currently viewing + if (convId === conversationsStore.activeConversation?.id) { + this.host.processing.setActiveConversation(convId); + } + + const unlock = () => { + this.host.setChatLoading(convId, false); + this.host.clearChatStreaming(convId); + }; + // fetch the replay stream from byte 0, rebuild the assistant message from scratch. + // resolve the server side identity, fall back to streamIdentity when the caller does not + // pass a streamId. probeServerStream returns the full id (with ::model suffix when present) + const id = streamId || streamIdentity(convId, modelsStore.selectedModelName); + + let response: Response; + + try { + response = await ChatService.fetchStreamReplay(id); + } catch (e) { + console.error(`attachServerStream replay failed for conv ${convId}:`, e); + unlock(); + + return; + } + + // load the target conversation messages by id, not via the active store. when multiple + // attaches run in parallel the active store may reflect another conv and writing through + // its index mixes content across convs (CoT flicker, message bleed). by going through the + // DB we stay isolated, and only mirror into the active store when the attached conv is + // the one currently displayed + let messages: DatabaseMessage[]; + + try { + messages = await DatabaseService.getConversationMessages(convId); + } catch (e) { + console.error('attachServerStream load messages failed:', e); + unlock(); + + return; + } + + // locate the slot to splice into, create a placeholder assistant message if there is none. + // we use the conv-scoped findLastAssistantIdx helpers, they only depend on the array + let targetIdx = this.findLastAssistantIdx(messages); + + if (targetIdx === -1) { + const lastUserIdx = this.findLastUserIdx(messages); + + if (lastUserIdx === -1) { + console.warn( + `attachServerStream: conv ${convId} has no user or assistant message, cannot splice` + ); + unlock(); + + return; + } + + try { + const placeholder = await DatabaseService.createMessageBranch( + { + children: [], + content: '', + convId, + parent: messages[lastUserIdx].id, + role: MessageRole.ASSISTANT, + timestamp: Date.now(), + toolCalls: '', + type: MessageType.TEXT + } as Omit, + messages[lastUserIdx].id + ); + + messages = [...messages, placeholder]; + targetIdx = messages.length - 1; + + // only push into the active store when this conv is the one displayed right now + if (convId === conversationsStore.activeConversation?.id) { + conversationsStore.addMessageToActive(placeholder); + } + } catch (e) { + console.error('attachServerStream placeholder creation failed:', e); + unlock(); + + return; + } + } + + if (targetIdx === -1) { + unlock(); + + return; + } + + const targetMessage = messages[targetIdx]; + const targetMessageId = targetMessage.id; + // when the assistant slot already has content, the running session is a continue or + // another append flow and its buffer holds only the appended deltas. preserve the prefix + // and let the replay add to it. when the slot is empty the session buffer holds the whole + // message so we wipe and rebuild from byte 0 + const existingContent = targetMessage.content ?? ''; + const existingReasoning = targetMessage.reasoningContent ?? ''; + const isAppendMode = existingContent.length > 0; + // helper: write to the active store only when the attached conv is currently displayed. + // the lookup by message id is robust to reordering of activeMessages, two parallel attaches + // can no longer step on each other's indices + const writeActive = (updates: Partial) => { + if (convId !== conversationsStore.activeConversation?.id) { + return; + } + + const liveIdx = conversationsStore.findMessageIndex(targetMessageId); + + if (liveIdx === -1) return; + + conversationsStore.updateMessageAtIndex(liveIdx, updates); + }; + + if (!isAppendMode) { + writeActive({ content: '', reasoningContent: undefined }); + } + + // extract the model suffix, the resume calls in handleStreamResponse must reuse the model + // the session was tagged with, not the live dropdown + const sepIdx = id.indexOf(CONVERSATION_ID_SEPARATOR); + const attachedModel: string | null = sepIdx === -1 ? null : id.slice(sepIdx + 2); + + this.host.setChatStreaming(convId, existingContent, targetMessageId, attachedModel); + const abortController = this.host.getOrCreateAbortController(convId); + + let streamedContent = ''; + let streamedReasoningContent = ''; + + const cleanup = () => { + unlock(); + this.host.processing.setState(convId, null); + }; + + try { + await ChatService.handleStreamResponse( + response, + (chunk: string) => { + streamedContent += chunk; + const displayed = isAppendMode ? existingContent + streamedContent : streamedContent; + + writeActive({ content: displayed }); + this.host.setChatStreaming(convId, displayed, targetMessageId); + }, + async ( + finalContent?: string, + reasoningContent?: string, + timings?: ChatMessageTimings, + toolCalls?: string + ) => { + const streamed = streamedContent || finalContent || ''; + const streamedR = streamedReasoningContent || reasoningContent || ''; + const content = isAppendMode ? existingContent + streamed : streamed; + const reasoning = isAppendMode ? existingReasoning + streamedR : streamedR; + + // the DB write is the source of truth, mirror to the active store only when + // the conv is currently displayed + await DatabaseService.updateMessage(targetMessageId, { + content, + reasoningContent: reasoning || undefined, + timings, + toolCalls: toolCalls || '' + }); + writeActive({ + content, + reasoningContent: reasoning || undefined, + timings + }); + cleanup(); + }, + (err: Error) => { + console.error('attachServerStream pipe error:', err); + cleanup(); + }, + (chunk: string) => { + streamedReasoningContent += chunk; + const displayed = isAppendMode + ? existingReasoning + streamedReasoningContent + : streamedReasoningContent; + + writeActive({ reasoningContent: displayed }); + }, + undefined, + undefined, + undefined, + undefined, + convId, + abortController.signal, + (connState: StreamConnectionState) => { + if (convId === conversationsStore.activeConversation?.id) { + this.host.streamConnectionState = connState; + } + }, + attachedModel + ); + } catch (e) { + console.error('attachServerStream pipe crashed:', e); + cleanup(); + } + } + + private findLastAssistantIdx(messages: DatabaseMessage[]): number { + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role === MessageRole.ASSISTANT) return i; + } + + return -1; + } + + private findLastUserIdx(messages: DatabaseMessage[]): number { + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].role === MessageRole.USER) return i; + } + + return -1; + } + + /** + * Server side stream discovery, split in three pieces: + * + * probeServerStream(convId) -> hits POST /v1/streams/lookup with the conv id, returns the session to attach + * to or null. Pure read, no side effect, no UI lock. Safe to fire in parallel with anything. + * + * attachServerStream(convId) -> flips the spinner immediately, fetches the replay stream + * from byte 0, finds the assistant slot to splice into (creates a placeholder if the conv has + * no assistant message yet, for cross device or fresh local DB cases), and pipes the SSE bytes + * into the message via handleStreamResponse. + * + * discoverActiveStream(convId) -> probe + attach in one call. Used by callers that do not need + * to overlap the probe with other async work. + * + * The chat page in +page.svelte calls discoverActiveStream once the conversation is active + * (immediately if it already is, after loadConversation settles otherwise), and re-runs it on + * visibilitychange. Attaching only after the conversation is loaded gives the earliest + * possible time to spinner and avoids racing against an empty activeMessages array. + */ + private async probeServerStream(convId: string): Promise { + if (!convId) return null; + + let sessions: ApiStreamSession[]; + + try { + sessions = await ChatService.lookupStreamSessions([convId]); + } catch (e) { + console.warn(`probeServerStream failed for conv ${convId}:`, e); + + return null; + } + + return ChatService.selectActiveStream(sessions); + } +} diff --git a/tools/ui/src/lib/stores/conversations.svelte.ts b/tools/ui/src/lib/stores/conversations/index.svelte.ts similarity index 61% rename from tools/ui/src/lib/stores/conversations.svelte.ts rename to tools/ui/src/lib/stores/conversations/index.svelte.ts index d2184b359355..7d6dc326c48a 100644 --- a/tools/ui/src/lib/stores/conversations.svelte.ts +++ b/tools/ui/src/lib/stores/conversations/index.svelte.ts @@ -1,135 +1,67 @@ /** - * conversationsStore - Reactive State Store for Conversations + * conversationsStore - Conversation lifecycle, persistence and navigation * - * Manages conversation lifecycle, persistence, navigation, and MCP server overrides. - * - * **Architecture & Relationships:** - * - **DatabaseService**: Stateless IndexedDB layer - * - **conversationsStore** (this): Reactive state + business logic - * - **chatStore**: Chat-specific state (streaming, loading) - * - * **Key Responsibilities:** - * - Conversation CRUD (create, load, delete) - * - Message management and tree navigation - * - MCP server per-chat overrides - * - Import/Export functionality - * - Title management with confirmation - * - * @see DatabaseService in services/database.ts for IndexedDB operations + * Owns conversation CRUD, message tree navigation, import/export and title + * management, persisted through DatabaseService. Per-chat options (MCP + * overrides, reasoning effort, cwd) live in ConversationPreferences, + * composed as {@link ConversationsStore.preferences}. */ import { browser } from '$app/environment'; import { goto } from '$app/navigation'; -import { REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY, ROUTES } from '$lib/constants'; -import { MessageRole, ReasoningEffort } from '$lib/enums'; +import { ROUTES } from '$lib/constants'; +import { MessageRole } from '$lib/enums'; import { ConversationTransferService } from '$lib/services/conversation-transfer.service'; import { DatabaseService } from '$lib/services/database.service'; import { MigrationService } from '$lib/services/migration.service'; import { RouterService } from '$lib/services/router.service'; // direct imports between stores, not via the barrel, to avoid circular deps -import { mcpStore } from '$lib/stores/mcp.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; -import type { McpServerOverride } from '$lib/types/database'; +import { + ConversationPreferences, + type ConversationsPreferencesHost +} from '$lib/stores/conversations/preferences.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { filterByLeafNodeId, findLeafNode, generateConversationTitle } from '$lib/utils'; import { SvelteSet } from 'svelte/reactivity'; import { toast } from 'svelte-sonner'; -class ConversationsStore { - /** - * - * - * State - * - * - */ - - /** List of all conversations */ - conversations = $state([]); - +class ConversationsStore implements ConversationsPreferencesHost { /** Currently active conversation */ activeConversation = $state(null); /** Messages in the active conversation (filtered by currNode path) */ activeMessages = $state([]); + /** List of all conversations */ + conversations = $state([]); + /** Whether the store has been initialized */ isInitialized = $state(false); - /** Global (non-conversation-specific) reasoning effort default */ - pendingReasoningEffort = $state(ConversationsStore.loadReasoningEffortDefault()); + /** Per-chat options (MCP overrides, reasoning effort, cwd), composed here. */ + private _preferences = new ConversationPreferences(this); /** - * Working directory picked on the empty new-chat screen, before any - * conversation exists. Consumed by `chatStore.sendMessage()`, which - * records it into chat history as a synthetic message on first send. - * Cleared by `loadConversation` and `clearActiveConversation` so a - * stale pick can't bleed onto an unrelated chat. + * Listeners notified with the ids of conversations that were deleted. + * Lets dependent stores (e.g. agenticStore) drop per-conversation state + * without introducing a circular import back into this store. */ - pendingCwd = $state(null); - - /** Load reasoning effort default from localStorage, DEFAULT defers to the server */ - private static loadReasoningEffortDefault(): ReasoningEffort { - if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.DEFAULT; - - try { - const raw = localStorage.getItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY); - - return (raw as ReasoningEffort) || ReasoningEffort.DEFAULT; - } catch { - return ReasoningEffort.DEFAULT; - } - } - - /** Persist reasoning effort default to localStorage */ - private saveReasoningEffortDefaults(): void { - if (typeof globalThis.localStorage === 'undefined') return; - - localStorage.setItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY, this.pendingReasoningEffort); - } + private conversationDeletionListeners = new Set<(convIds: string[]) => void>(); /** In-flight init run; shared by concurrent callers, reset on failure to allow retry */ private initPromise: Promise | null = null; /** - * - * - * Lifecycle - * - * - */ - - /** - * Initialize the store by loading conversations from database. - * Safe to call multiple times: concurrent callers share a single run, - * and a failed run can be retried by calling again. + * Memo of the last findMessageIndex() lookup. Streaming calls it once per + * chunk for the same message, so a validated cache hit keeps that O(1) + * instead of a linear scan of activeMessages on every token. */ - init(): Promise { - if (!browser) return Promise.resolve(); - - if (this.initPromise) return this.initPromise; - - this.initPromise = (async () => { - try { - await MigrationService.runAllMigrations(); - await this.loadConversations(); - this.isInitialized = true; - } catch (error) { - console.error('Failed to initialize conversations:', error); - this.initPromise = null; - } - })(); + private lastMessageIndex: { id: string; index: number } | null = null; - return this.initPromise; + get preferences() { + return this._preferences; } - /** - * - * - * Message Array Operations - * - * - */ - /** * Adds a message to the active messages array */ @@ -138,67 +70,172 @@ class ConversationsStore { } /** - * Updates a message at a specific index in active messages + * Applies a field update to a conversation row, mirroring it into both the + * conversations list and the active conversation when it is the target. + * Shared by the rename/pin/preferences flows so no caller can forget to + * mirror one side. */ - updateMessageAtIndex(index: number, updates: Partial): void { - const message = index === -1 ? undefined : this.activeMessages[index]; + applyConversationUpdate(id: string, updates: Partial): void { + const convIndex = this.conversations.findIndex((c) => c.id === id); - if (!message) return; - - // Assign field by field rather than replacing the object. Replacing it - // changes the array slot, which invalidates every consumer that merely - // walks the list - notably ChatMessages.displayMessages, which rebuilds - // entries for every message in the conversation. Deep $state proxies make - // per-field writes fine-grained, so only readers of the changed field wake. - const target = message as unknown as Record; + if (convIndex !== -1) { + const target = this.conversations[convIndex] as unknown as Record; - for (const [key, value] of Object.entries(updates)) { - if (target[key] !== value) { - target[key] = value; + for (const [key, value] of Object.entries(updates)) { + if (target[key] !== value) target[key] = value; } } + + if (this.activeConversation?.id === id) { + this.activeConversation = { ...this.activeConversation, ...updates }; + } } /** - * Finds the index of a message in active messages + * Derives a conversation title from its first message content and applies + * it, honoring the title-generation setting. Shared by every flow that + * edits or creates the first user message. */ - findMessageIndex(messageId: string): number { - return this.activeMessages.findIndex((m) => m.id === messageId); + async applyTitleFromContent(convId: string, content: string): Promise { + await this.updateConversationName( + convId, + generateConversationTitle(content, Boolean(settingsStore.config.titleGenerationUseFirstLine)) + ); } /** - * Removes messages from active messages starting at an index + * Deletes multiple conversations in sequence. + * Mirrors deleteConversation() per-id; navigates to NEW_CHAT only if the + * currently-open chat was among the deleted ones. + * @param convIds - Conversation IDs to delete */ - sliceActiveMessages(startIndex: number): void { - this.activeMessages = this.activeMessages.slice(0, startIndex); + async bulkDeleteConversations(convIds: string[]): Promise { + if (convIds.length === 0) return; + + try { + const idsToRemove = new SvelteSet(convIds); + // Collect all descendants recursively so the local cache stays consistent + // even when deleteWithForks is omitted. + const queue = [...convIds]; + + while (queue.length > 0) { + const parentId = queue.pop()!; + + for (const c of this.conversations) { + if (c.forkedFromConversationId === parentId && !idsToRemove.has(c.id)) { + idsToRemove.add(c.id); + queue.push(c.id); + } + } + } + + const activeWasDeleted = + this.activeConversation !== null && idsToRemove.has(this.activeConversation.id); + + await DatabaseService.bulkDeleteConversations([...idsToRemove]); + + this.conversations = this.conversations.filter((c) => !idsToRemove.has(c.id)); + this.notifyConversationsDeleted([...idsToRemove]); + + if (activeWasDeleted) { + this.clearActiveConversation(); + await goto(ROUTES.NEW_CHAT); + } + + toast.success( + idsToRemove.size === 1 + ? 'Conversation deleted' + : `${idsToRemove.size} conversations deleted` + ); + } catch (error) { + console.error('Failed to bulk delete conversations:', error); + toast.error('Failed to delete conversations'); + } } /** - * Removes a message from active messages by index + * Bundles the given conversations into a single zip archive and triggers a + * browser download (one JSONL file per conversation). + * @param convIds - Conversation IDs to export */ - removeMessageAtIndex(index: number): DatabaseMessage | undefined { - if (index !== -1) { - return this.activeMessages.splice(index, 1)[0]; - } + async bulkExportConversations(convIds: string[]): Promise { + if (convIds.length === 0) return; - return undefined; + try { + const fetched = await DatabaseService.getConversationsWithMessages(convIds); + const activeId = this.activeConversation?.id; + const overridden = fetched.get(activeId ?? ''); + + if (overridden && activeId) { + overridden.conv = { ...this.activeConversation! }; + } + + const exported = [...fetched.values()]; + + if (exported.length === 0) { + toast.error('No conversations to export'); + + return; + } + + ConversationTransferService.downloadConversationsArchive(exported); + + toast.success( + exported.length === 1 + ? 'Conversation exported' + : `${exported.length} conversations exported` + ); + } catch (error) { + console.error('Failed to bulk export conversations:', error); + toast.error('Failed to export conversations'); + } } /** - * - * - * Conversation CRUD - * - * + * Toggles the pinned state of each conversation individually. + * Mixed-pin selections are intentionally not normalised here; the bulk + * action UI surfaces them as a disabled mixed-state instead. + * @param convIds - Conversation IDs to toggle */ + async bulkToggleConversationPin(convIds: string[]): Promise { + if (convIds.length === 0) return; + + try { + const updates = await DatabaseService.bulkToggleConversationPins(convIds); + const activeId = this.activeConversation?.id; + + if (activeId && updates.has(activeId)) { + this.activeConversation = { + ...this.activeConversation!, + pinned: updates.get(activeId)! + }; + } + + for (let i = 0; i < this.conversations.length; i++) { + const newPinned = updates.get(this.conversations[i].id); + + if (newPinned !== undefined) this.conversations[i].pinned = newPinned; + } + + toast.success( + convIds.length === 1 + ? 'Conversation pin toggled' + : `Updated pin state for ${convIds.length} conversations` + ); + } catch (error) { + console.error('Failed to bulk toggle pin:', error); + toast.error('Failed to update pin state'); + } + } /** - * Loads all conversations from the database + * Clears the active conversation and messages. */ - async loadConversations(): Promise { - const conversations = await DatabaseService.getAllConversations(); - - this.conversations = conversations; + clearActiveConversation(): void { + this.activeConversation = null; + this.activeMessages = []; + // reload defaults so new chats inherit persisted state + this.preferences.resetPending(); } /** @@ -208,17 +245,15 @@ class ConversationsStore { */ async createConversation(name?: string): Promise { const conversationName = name || `Chat ${new Date().toLocaleString()}`; - // No MCP override list is seeded: getAllMcpServerOverrides resolves - // servers without a per-conversation override to `mcpServers[i].enabled`, - // and only explicit toggles are stored on the conversation. - // Working directory picked on the new-chat screen gets threaded in - // here too, then cleared so it doesn't bleed onto subsequent new chats. + // Working directory and reasoning effort picked on the new-chat screen + // get threaded into the new conversation here, then cleared so they + // don't bleed onto subsequent new chats. const conversation = await DatabaseService.createConversation(conversationName, { - cwd: this.pendingCwd ?? undefined, - reasoningEffort: this.pendingReasoningEffort + cwd: this.preferences.pendingCwd ?? undefined, + reasoningEffort: this.preferences.pendingReasoningEffort }); - this.pendingCwd = null; + this.preferences.pendingCwd = null; this.conversations = [conversation, ...this.conversations]; this.activeConversation = conversation; @@ -230,58 +265,28 @@ class ConversationsStore { } /** - * Loads a specific conversation and its messages - * @param convId - The conversation ID to load - * @returns True if conversation was loaded successfully + * Deletes all conversations and their messages */ - async loadConversation(convId: string): Promise { + async deleteAll(): Promise { try { - const conversation = await DatabaseService.getConversation(convId); - - if (!conversation) { - return false; - } - - // Drop any cwd the user drafted on the empty new-chat screen - - // it doesn't belong to this conversation. - this.pendingCwd = null; - - this.activeConversation = conversation; + const allConversations = await DatabaseService.getAllConversations(); + const allIds = allConversations.map((c) => c.id); - if (conversation.currNode) { - const allMessages = await DatabaseService.getConversationMessages(convId); - const filteredMessages = filterByLeafNodeId( - allMessages, - conversation.currNode, - false - ) as DatabaseMessage[]; + await DatabaseService.bulkDeleteConversations(allIds); - this.activeMessages = filteredMessages; - } else { - const messages = await DatabaseService.getConversationMessages(convId); + this.clearActiveConversation(); + this.conversations = []; + this.notifyConversationsDeleted(allIds); - this.activeMessages = messages; - } + toast.success('All conversations deleted'); - return true; + await goto(ROUTES.NEW_CHAT); } catch (error) { - console.error('Failed to load conversation:', error); - - return false; + console.error('Failed to delete all conversations:', error); + toast.error('Failed to delete conversations'); } } - /** - * Clears the active conversation and messages. - */ - clearActiveConversation(): void { - this.activeConversation = null; - this.activeMessages = []; - // reload defaults so new chats inherit persisted state - this.pendingReasoningEffort = ConversationsStore.loadReasoningEffortDefault(); - this.pendingCwd = null; - } - /** * Deletes a conversation and all its messages * @param convId - The conversation ID to delete @@ -311,6 +316,8 @@ class ConversationsStore { this.clearActiveConversation(); await goto(ROUTES.NEW_CHAT); } + + this.notifyConversationsDeleted([...idsToRemove]); } else { // Reparent direct children to deleted conv's parent (or promote to top-level) const deletedConv = this.conversations.find((c) => c.id === convId); @@ -328,6 +335,8 @@ class ConversationsStore { this.clearActiveConversation(); await goto(ROUTES.NEW_CHAT); } + + this.notifyConversationsDeleted([convId]); } } catch (error) { console.error('Failed to delete conversation:', error); @@ -335,178 +344,87 @@ class ConversationsStore { } /** - * Deletes all conversations and their messages + * Downloads a single conversation as a JSONL file, serializing the full message tree. + * @param convId - The conversation ID to download */ - async deleteAll(): Promise { - try { - const allConversations = await DatabaseService.getAllConversations(); - - await DatabaseService.bulkDeleteConversations(allConversations.map((c) => c.id)); + async downloadConversation(convId: string): Promise { + const conversation = + this.activeConversation?.id === convId + ? this.activeConversation + : await DatabaseService.getConversation(convId); - this.clearActiveConversation(); - this.conversations = []; + if (!conversation) return; - toast.success('All conversations deleted'); + const messages = await DatabaseService.getConversationMessages(convId); - await goto(ROUTES.NEW_CHAT); - } catch (error) { - console.error('Failed to delete all conversations:', error); - toast.error('Failed to delete conversations'); - } + ConversationTransferService.downloadConversationFile({ conv: conversation, messages }); } /** - * Deletes multiple conversations in sequence. - * Mirrors deleteConversation() per-id; navigates to NEW_CHAT only if the - * currently-open chat was among the deleted ones. - * @param convIds - Conversation IDs to delete + * Finds the index of a message in active messages. + * + * The last lookup is memoized and reused when it still validates against + * the current array (same id at the same position), which covers the + * streaming hot path where the same message is looked up on every chunk + * while the array itself only mutates by field. Any structural change + * (splice, reassignment, reordering) fails validation and falls back to a + * full scan. */ - async bulkDeleteConversations(convIds: string[]): Promise { - if (convIds.length === 0) return; - - try { - const idsToRemove = new SvelteSet(convIds); - // Collect all descendants recursively so the local cache stays consistent - // even when deleteWithForks is omitted. - const queue = [...convIds]; - - while (queue.length > 0) { - const parentId = queue.pop()!; - - for (const c of this.conversations) { - if (c.forkedFromConversationId === parentId && !idsToRemove.has(c.id)) { - idsToRemove.add(c.id); - queue.push(c.id); - } - } - } - - const activeWasDeleted = - this.activeConversation !== null && idsToRemove.has(this.activeConversation.id); - - await DatabaseService.bulkDeleteConversations([...idsToRemove]); - - this.conversations = this.conversations.filter((c) => !idsToRemove.has(c.id)); - - if (activeWasDeleted) { - this.clearActiveConversation(); - await goto(ROUTES.NEW_CHAT); - } - - toast.success( - idsToRemove.size === 1 - ? 'Conversation deleted' - : `${idsToRemove.size} conversations deleted` - ); - } catch (error) { - console.error('Failed to bulk delete conversations:', error); - toast.error('Failed to delete conversations'); + findMessageIndex(messageId: string): number { + const last = this.lastMessageIndex; + const messages = this.activeMessages; + + if ( + last && + last.id === messageId && + last.index >= 0 && + last.index < messages.length && + messages[last.index]?.id === messageId + ) { + return last.index; } - } - - /** - * Toggles the pinned state of each conversation individually. - * Mixed-pin selections are intentionally not normalised here; the bulk - * action UI surfaces them as a disabled mixed-state instead. - * @param convIds - Conversation IDs to toggle - */ - async bulkToggleConversationPin(convIds: string[]): Promise { - if (convIds.length === 0) return; - try { - const updates = await DatabaseService.bulkToggleConversationPins(convIds); - const activeId = this.activeConversation?.id; + const index = messages.findIndex((m) => m.id === messageId); - if (activeId && updates.has(activeId)) { - this.activeConversation = { - ...this.activeConversation!, - pinned: updates.get(activeId)! - }; - } + this.lastMessageIndex = { id: messageId, index }; - for (let i = 0; i < this.conversations.length; i++) { - const newPinned = updates.get(this.conversations[i].id); - - if (newPinned !== undefined) this.conversations[i].pinned = newPinned; - } - - toast.success( - convIds.length === 1 - ? 'Conversation pin toggled' - : `Updated pin state for ${convIds.length} conversations` - ); - } catch (error) { - console.error('Failed to bulk toggle pin:', error); - toast.error('Failed to update pin state'); - } + return index; } /** - * Bundles the given conversations into a single zip archive and triggers a - * browser download (one JSONL file per conversation). - * @param convIds - Conversation IDs to export + * Forks a conversation at a specific message, creating a new conversation + * containing messages from root up to the target message, then navigates to it. + * + * @param messageId - The message ID to fork at + * @param options - Fork options (name and whether to include attachments) + * @returns The new conversation ID, or null if fork failed */ - async bulkExportConversations(convIds: string[]): Promise { - if (convIds.length === 0) return; + async forkConversation( + messageId: string, + options: { name: string; includeAttachments: boolean } + ): Promise { + if (!this.activeConversation) return null; try { - const fetched = await DatabaseService.getConversationsWithMessages(convIds); - const activeId = this.activeConversation?.id; - const overridden = fetched.get(activeId ?? ''); - - if (overridden && activeId) { - overridden.conv = { ...this.activeConversation! }; - } - - const exported = [...fetched.values()]; - - if (exported.length === 0) { - toast.error('No conversations to export'); - - return; - } - - ConversationTransferService.downloadConversationsArchive(exported); - - toast.success( - exported.length === 1 - ? 'Conversation exported' - : `${exported.length} conversations exported` + const newConv = await DatabaseService.forkConversation( + this.activeConversation.id, + messageId, + options ); - } catch (error) { - console.error('Failed to bulk export conversations:', error); - toast.error('Failed to export conversations'); - } - } - /** - * - * - * Message Management - * - * - */ + this.conversations = [newConv, ...this.conversations]; - /** - * Refreshes active messages based on currNode after branch navigation. - */ - async refreshActiveMessages(): Promise { - if (!this.activeConversation) return; + await goto(RouterService.chat(newConv.id)); - const allMessages = await DatabaseService.getConversationMessages(this.activeConversation.id); + toast.success('Conversation forked'); - if (allMessages.length === 0) { - this.activeMessages = []; + return newConv.id; + } catch (error) { + console.error('Failed to fork conversation:', error); + toast.error('Failed to fork conversation'); - return; + return null; } - - const leafNodeId = - this.activeConversation.currNode || - allMessages.reduce((latest, msg) => (msg.timestamp > latest.timestamp ? msg : latest)).id; - const currentPath = filterByLeafNodeId(allMessages, leafNodeId, false) as DatabaseMessage[]; - - this.activeMessages = currentPath; } /** @@ -519,112 +437,95 @@ class ConversationsStore { } /** - * - * - * Title Management - * - * + * Imports conversations from provided data (without file picker) + * @param data - Array of conversation data with messages + * @returns The conversations written to the database and the ones skipped */ + async importConversationsData( + data: ExportedConversations + ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { + const result = await DatabaseService.importConversations(data); + + await this.loadConversations(); + + return result; + } /** - * Updates the name of a conversation. - * @param convId - The conversation ID to update - * @param name - The new name for the conversation + * Initialize the store by loading conversations from database. + * Safe to call multiple times: concurrent callers share a single run, + * and a failed run can be retried by calling again. */ - async updateConversationName(convId: string, name: string): Promise { - try { - await DatabaseService.updateConversation(convId, { name }); + initialize(): Promise { + if (!browser) return Promise.resolve(); - const convIndex = this.conversations.findIndex((c) => c.id === convId); + if (this.initPromise) return this.initPromise; - if (convIndex !== -1) { - this.conversations[convIndex].name = name; + this.initPromise = (async () => { + try { + await MigrationService.runAllMigrations(); + await this.loadConversations(); + this.isInitialized = true; + } catch (error) { + console.error('Failed to initialize conversations:', error); + this.initPromise = null; } + })(); - if (this.activeConversation?.id === convId) { - this.activeConversation = { ...this.activeConversation, name }; - } - } catch (error) { - console.error('Failed to update conversation name:', error); - } + return this.initPromise; } /** - * Toggles the pinned status of a conversation. - * @param convId - The conversation ID to toggle - * @returns The new pinned status + * Loads a specific conversation and its messages + * @param convId - The conversation ID to load + * @returns True if conversation was loaded successfully */ - async toggleConversationPin(convId: string): Promise { + async loadConversation(convId: string): Promise { try { - const newPinnedState = await DatabaseService.toggleConversationPin(convId); - const convIndex = this.conversations.findIndex((c) => c.id === convId); + const conversation = await DatabaseService.getConversation(convId); + + if (!conversation) { + return false; + } + + // Drop any cwd the user drafted on the empty new-chat screen - + // it doesn't belong to this conversation. + this.preferences.pendingCwd = null; + + this.activeConversation = conversation; + + if (conversation.currNode) { + const allMessages = await DatabaseService.getConversationMessages(convId); + const filteredMessages = filterByLeafNodeId( + allMessages, + conversation.currNode, + false + ) as DatabaseMessage[]; - if (convIndex !== -1) { - this.conversations[convIndex].pinned = newPinnedState; - } + this.activeMessages = filteredMessages; + } else { + const messages = await DatabaseService.getConversationMessages(convId); - if (this.activeConversation?.id === convId) { - this.activeConversation = { ...this.activeConversation, pinned: newPinnedState }; + this.activeMessages = messages; } - return newPinnedState; + return true; } catch (error) { - console.error('Failed to toggle conversation pin:', error); + console.error('Failed to load conversation:', error); return false; } } /** - * Marks a conversation as recently active: stamps lastModified (persisted) - * and moves it to the top of the list. Only message-activity flows call - * this; metadata updates (rename, pin, settings) do not. - * - * @param convId - Conversation that produced the activity, defaults to the active one - */ - updateConversationTimestamp(convId?: string): void { - const targetId = convId ?? this.activeConversation?.id; - - if (!targetId) return; - - const now = Date.now(); - const chatIndex = this.conversations.findIndex((c) => c.id === targetId); - - if (chatIndex !== -1) { - this.conversations[chatIndex].lastModified = now; - const updatedConv = this.conversations.splice(chatIndex, 1)[0]; - - this.conversations = [updatedConv, ...this.conversations]; - } - - if (this.activeConversation?.id === targetId) { - this.activeConversation = { ...this.activeConversation, lastModified: now }; - } - - DatabaseService.updateConversation(targetId, { lastModified: now }).catch((error) => - console.error('Failed to update conversation timestamp:', error) - ); - } - - /** - * Updates the current node of the active conversation - * @param nodeId - The new current node ID + * Loads all conversations from the database */ - async updateCurrentNode(nodeId: string): Promise { - if (!this.activeConversation) return; + async loadConversations(): Promise { + const conversations = await DatabaseService.getAllConversations(); - await DatabaseService.updateCurrentNode(this.activeConversation.id, nodeId); - this.activeConversation = { ...this.activeConversation, currNode: nodeId }; + this.conversations = conversations; } - /** - * - * - * Branch Navigation - * - * - */ - /** * Navigates to a specific sibling branch by updating currNode and refreshing messages. * @param siblingId - The sibling message ID to navigate to @@ -655,279 +556,124 @@ class ConversationsStore { newFirstUserMessage.id !== currentFirstUserMessage.id || newFirstUserMessage.content.trim() !== currentFirstUserMessage.content.trim()) ) { - await this.updateConversationName( - this.activeConversation.id, - generateConversationTitle( - newFirstUserMessage.content, - Boolean(settingsStore.config.titleGenerationUseFirstLine) - ) - ); + await this.applyTitleFromContent(this.activeConversation.id, newFirstUserMessage.content); } } } /** - * - * - * MCP Server Overrides - * - * - */ - - /** - * Resolve the default enabled value for a server: its own `enabled` - * flag in `mcpServers`, so the global on/off state lives in one place. - */ - #getDefaultOverride(serverId: string): McpServerOverride | undefined { - const server = mcpStore.getServers().find((s) => s.id === serverId); - - if (!server) return undefined; - - return { enabled: server.enabled, serverId }; - } - - /** - * Gets the effective MCP server override for a specific server. - * A per-conversation override wins when present; a server without one - * resolves to its `mcpServers[i].enabled` default. - * @param serverId - The server ID to check - * @returns The effective override, undefined if no matching server - */ - getMcpServerOverride(serverId: string): McpServerOverride | undefined { - const override = this.activeConversation?.mcpServerOverrides?.find( - (o: McpServerOverride) => o.serverId === serverId - ); - - if (override) return override; - - return this.#getDefaultOverride(serverId); - } - - /** - * Gets the effective override list for the current conversation: - * one entry per configured server, resolved per server. The stored - * per-conversation list is sparse and only holds explicit toggles. + * Registers a listener invoked with the ids of deleted conversations. + * Returns an unsubscribe function. */ - getAllMcpServerOverrides(): McpServerOverride[] { - const overrides = this.activeConversation?.mcpServerOverrides; + onConversationsDeleted(listener: (convIds: string[]) => void): () => void { + this.conversationDeletionListeners.add(listener); - return mcpStore.getServers().map((s) => { - const override = overrides?.find((o: McpServerOverride) => o.serverId === s.id); - - return { enabled: override?.enabled ?? s.enabled, serverId: s.id }; - }); + return () => this.conversationDeletionListeners.delete(listener); } /** - * Checks if an MCP server is enabled for the active conversation. - * @param serverId - The server ID to check - * @returns True if server is enabled for this conversation + * Refreshes active messages based on currNode after branch navigation. */ - isMcpServerEnabledForChat(serverId: string): boolean { - const override = this.getMcpServerOverride(serverId); + async refreshActiveMessages(): Promise { + if (!this.activeConversation) return; - return override?.enabled ?? false; - } + const allMessages = await DatabaseService.getConversationMessages(this.activeConversation.id); - /** - * Sets or removes MCP server override for the active conversation. - * If no conversation exists, persists `enabled` onto `mcpServers[i].enabled` - * (the single source of truth for new-chat defaults). - * @param serverId - The server ID to override - * @param enabled - The enabled state, or undefined to remove per-conversation override - */ - async setMcpServerOverride(serverId: string, enabled: boolean | undefined): Promise { - if (!this.activeConversation) { - if (enabled !== undefined) { - mcpStore.updateServer(serverId, { enabled }); - } + if (allMessages.length === 0) { + this.activeMessages = []; return; } - // Clone to plain objects to avoid Proxy serialization issues with IndexedDB - const currentOverrides = (this.activeConversation.mcpServerOverrides || []).map( - (o: McpServerOverride) => ({ - enabled: o.enabled, - serverId: o.serverId - }) - ); - - let newOverrides: McpServerOverride[]; - - if (enabled === undefined) { - newOverrides = currentOverrides.filter((o: McpServerOverride) => o.serverId !== serverId); - } else { - const existingIndex = currentOverrides.findIndex( - (o: McpServerOverride) => o.serverId === serverId - ); - - if (existingIndex >= 0) { - newOverrides = [...currentOverrides]; - newOverrides[existingIndex] = { enabled, serverId }; - } else { - newOverrides = [...currentOverrides, { enabled, serverId }]; - } - } - - await DatabaseService.updateConversation(this.activeConversation.id, { - mcpServerOverrides: newOverrides.length > 0 ? newOverrides : undefined - }); - - this.activeConversation = { - ...this.activeConversation, - mcpServerOverrides: newOverrides.length > 0 ? newOverrides : undefined - }; - - const convIndex = this.conversations.findIndex((c) => c.id === this.activeConversation!.id); + const leafNodeId = + this.activeConversation.currNode || + allMessages.reduce((latest, msg) => (msg.timestamp > latest.timestamp ? msg : latest)).id; + const currentPath = filterByLeafNodeId(allMessages, leafNodeId, false) as DatabaseMessage[]; - if (convIndex !== -1) { - this.conversations[convIndex].mcpServerOverrides = - newOverrides.length > 0 ? newOverrides : undefined; - } + this.activeMessages = currentPath; } /** - * Toggles MCP server enabled state for the active conversation. - * @param serverId - The server ID to toggle + * Removes a message from active messages by index */ - async toggleMcpServerForChat(serverId: string): Promise { - const currentEnabled = this.isMcpServerEnabledForChat(serverId); - - await this.setMcpServerOverride(serverId, !currentEnabled); - } + removeMessageAtIndex(index: number): DatabaseMessage | undefined { + if (index !== -1) { + return this.activeMessages.splice(index, 1)[0]; + } - /** - * Removes MCP server override for the active conversation. - * @param serverId - The server ID to remove override for - */ - async removeMcpServerOverride(serverId: string): Promise { - await this.setMcpServerOverride(serverId, undefined); + return undefined; } /** - * Gets the effective reasoning effort for the active conversation. - * Returns the conversation override if set, otherwise the global default. - * DEFAULT means no override is sent and the server decides. + * Removes messages from active messages starting at an index */ - getReasoningEffort(): ReasoningEffort { - if (this.activeConversation) { - if (this.activeConversation.reasoningEffort !== undefined) { - return this.activeConversation.reasoningEffort; - } - - // conversations created before the tri-state store an explicit - // opt-out only as thinkingEnabled = false - if (this.activeConversation.thinkingEnabled === false) { - return ReasoningEffort.OFF; - } - } - - return this.pendingReasoningEffort; + sliceActiveMessages(startIndex: number): void { + this.activeMessages = this.activeMessages.slice(0, startIndex); } /** - * Sets the reasoning effort for the active conversation. - * If no conversation exists, stores the global default. - * @param effort - The effort level ('default' | 'off' | 'low' | 'medium' | 'high' | 'max') + * Toggles the pinned status of a conversation. + * @param convId - The conversation ID to toggle + * @returns The new pinned status */ - async setReasoningEffort(effort: ReasoningEffort): Promise { - if (!this.activeConversation) { - this.pendingReasoningEffort = effort; - this.saveReasoningEffortDefaults(); - - return; - } - - this.activeConversation = { - ...this.activeConversation, - reasoningEffort: effort - }; + async toggleConversationPin(convId: string): Promise { + try { + const newPinnedState = await DatabaseService.toggleConversationPin(convId); - await DatabaseService.updateConversation(this.activeConversation.id, { - reasoningEffort: effort - }); + this.applyConversationUpdate(convId, { pinned: newPinnedState }); - const convIndex = this.conversations.findIndex((c) => c.id === this.activeConversation!.id); + return newPinnedState; + } catch (error) { + console.error('Failed to toggle conversation pin:', error); - if (convIndex !== -1) { - this.conversations[convIndex].reasoningEffort = effort; + return false; } } /** - * Sets the working directory for the active conversation. Pass `null` or - * an empty string to clear it, which restores the picker's empty state. - * - * On the empty new-chat screen (no active conversation yet), the value - * is buffered into `pendingCwd` so the user can pick before - * sending the first message; `createConversation()` consumes it. - * - * @param value - Absolute server-side path to the working directory, or null to clear + * Updates the name of a conversation. + * @param convId - The conversation ID to update + * @param name - The new name for the conversation */ - async setCwd(value: string | null): Promise { - const trimmed = value?.trim() || undefined; - - // No chat yet - buffer for the first chat the user creates. - if (!this.activeConversation) { - this.pendingCwd = trimmed ?? null; - - return; - } - - this.activeConversation = { - ...this.activeConversation, - cwd: trimmed - }; - - await DatabaseService.updateConversation(this.activeConversation.id, { - cwd: trimmed - }); - - const convIndex = this.conversations.findIndex((c) => c.id === this.activeConversation!.id); + async updateConversationName(convId: string, name: string): Promise { + try { + await DatabaseService.updateConversation(convId, { name }); - if (convIndex !== -1) { - this.conversations[convIndex].cwd = trimmed; - this.conversations = [...this.conversations]; + this.applyConversationUpdate(convId, { name }); + } catch (error) { + console.error('Failed to update conversation name:', error); } - - this.pendingCwd = null; } /** - * Forks a conversation at a specific message, creating a new conversation - * containing messages from root up to the target message, then navigates to it. + * Marks a conversation as recently active: stamps lastModified (persisted) + * and moves it to the top of the list. Only message-activity flows call + * this; metadata updates (rename, pin, settings) do not. * - * @param messageId - The message ID to fork at - * @param options - Fork options (name and whether to include attachments) - * @returns The new conversation ID, or null if fork failed + * @param convId - Conversation that produced the activity, defaults to the active one */ - async forkConversation( - messageId: string, - options: { name: string; includeAttachments: boolean } - ): Promise { - if (!this.activeConversation) return null; - - try { - const newConv = await DatabaseService.forkConversation( - this.activeConversation.id, - messageId, - options - ); + updateConversationTimestamp(convId?: string): void { + const targetId = convId ?? this.activeConversation?.id; - this.conversations = [newConv, ...this.conversations]; + if (!targetId) return; - await goto(RouterService.chat(newConv.id)); + const now = Date.now(); + const chatIndex = this.conversations.findIndex((c) => c.id === targetId); - toast.success('Conversation forked'); + if (chatIndex !== -1) { + this.conversations[chatIndex].lastModified = now; + const updatedConv = this.conversations.splice(chatIndex, 1)[0]; - return newConv.id; - } catch (error) { - console.error('Failed to fork conversation:', error); - toast.error('Failed to fork conversation'); + this.conversations = [updatedConv, ...this.conversations]; + } - return null; + if (this.activeConversation?.id === targetId) { + this.activeConversation = { ...this.activeConversation, lastModified: now }; } + + DatabaseService.updateConversation(targetId, { lastModified: now }).catch((error) => + console.error('Failed to update conversation timestamp:', error) + ); } /** @@ -939,35 +685,44 @@ class ConversationsStore { */ /** - * Downloads a single conversation as a JSONL file, serializing the full message tree. - * @param convId - The conversation ID to download + * Updates the current node of the active conversation + * @param nodeId - The new current node ID */ - async downloadConversation(convId: string): Promise { - const conversation = - this.activeConversation?.id === convId - ? this.activeConversation - : await DatabaseService.getConversation(convId); - - if (!conversation) return; - - const messages = await DatabaseService.getConversationMessages(convId); + async updateCurrentNode(nodeId: string): Promise { + if (!this.activeConversation) return; - ConversationTransferService.downloadConversationFile({ conv: conversation, messages }); + await DatabaseService.updateCurrentNode(this.activeConversation.id, nodeId); + this.activeConversation = { ...this.activeConversation, currNode: nodeId }; } /** - * Imports conversations from provided data (without file picker) - * @param data - Array of conversation data with messages - * @returns The conversations written to the database and the ones skipped + * Updates a message at a specific index in active messages */ - async importConversationsData( - data: ExportedConversations - ): Promise<{ imported: DatabaseConversation[]; skipped: DatabaseConversation[] }> { - const result = await DatabaseService.importConversations(data); + updateMessageAtIndex(index: number, updates: Partial): void { + const message = index === -1 ? undefined : this.activeMessages[index]; - await this.loadConversations(); + if (!message) return; - return result; + // Assign field by field rather than replacing the object. Replacing it + // changes the array slot, which invalidates every consumer that merely + // walks the list - notably ChatMessages.displayMessages, which rebuilds + // entries for every message in the conversation. Deep $state proxies make + // per-field writes fine-grained, so only readers of the changed field wake. + const target = message as unknown as Record; + + for (const [key, value] of Object.entries(updates)) { + if (target[key] !== value) { + target[key] = value; + } + } + } + + private notifyConversationsDeleted(convIds: string[]): void { + if (convIds.length === 0) return; + + for (const listener of this.conversationDeletionListeners) { + listener(convIds); + } } } diff --git a/tools/ui/src/lib/stores/conversations/preferences.svelte.ts b/tools/ui/src/lib/stores/conversations/preferences.svelte.ts new file mode 100644 index 000000000000..65a02344b69b --- /dev/null +++ b/tools/ui/src/lib/stores/conversations/preferences.svelte.ts @@ -0,0 +1,254 @@ +/** + * ConversationPreferences - Per-chat options with global fallback + * + * Owns the options that resolve per conversation: MCP server overrides, + * reasoning effort, and the working directory. Cwd and reasoning effort are + * buffered as pending state and threaded into the next created conversation + * by the host; MCP server overrides edit the sparse `mcpServerOverrides` + * list on the active row (new-chat toggles edit the server's global flag). + * Created and owned by conversationsStore; the host owns the conversation + * rows these options persist onto. + */ + +import { REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY } from '$lib/constants'; +import { ReasoningEffort } from '$lib/enums'; +import { DatabaseService } from '$lib/services/database.service'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { mcpStore } from '$lib/stores/mcp/index.svelte'; +import type { McpServerOverride } from '$lib/types/database'; + +/** Load reasoning effort default from localStorage, DEFAULT defers to the server */ +function loadReasoningEffortDefault(): ReasoningEffort { + if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.DEFAULT; + + try { + const raw = localStorage.getItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY); + + return (raw as ReasoningEffort) || ReasoningEffort.DEFAULT; + } catch { + return ReasoningEffort.DEFAULT; + } +} + +/** Persist reasoning effort default to localStorage */ +function saveReasoningEffortDefault(effort: ReasoningEffort): void { + if (typeof globalThis.localStorage === 'undefined') return; + + localStorage.setItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY, effort); +} + +/** + * The slice of conversationsStore the preferences read and write. Kept narrow + * on purpose so they cannot reach around the host's full surface; + * conversationsStore implements this structurally. + */ +export interface ConversationsPreferencesHost { + activeConversation: DatabaseConversation | null; + conversations: DatabaseConversation[]; + applyConversationUpdate(id: string, updates: Partial): void; +} + +export class ConversationPreferences { + /** + * Working directory picked on the empty new-chat screen, before any + * conversation exists. Consumed by `chatStore.sendMessage()`, which + * records it into chat history as a synthetic message on first send. + * Cleared by `loadConversation` and `clearActiveConversation` so a + * stale pick can't bleed onto an unrelated chat. + */ + pendingCwd = $state(null); + + /** Global (non-conversation-specific) reasoning effort default */ + pendingReasoningEffort = $state(loadReasoningEffortDefault()); + + constructor(private host: ConversationsPreferencesHost) {} + + /** + * Gets the effective override list for the current conversation: + * one entry per configured server, resolved per server. The stored + * per-conversation list is sparse and only holds explicit toggles. + */ + getAllMcpServerOverrides(): McpServerOverride[] { + const overrides = this.host.activeConversation?.mcpServerOverrides; + + return mcpStore.getServers().map((s) => { + const override = overrides?.find((o: McpServerOverride) => o.serverId === s.id); + + return { enabled: override?.enabled ?? s.enabled, serverId: s.id }; + }); + } + + /** + * Gets the effective MCP server override for a specific server. + * A per-conversation override wins when present; a server without one + * resolves to its `mcpServers[i].enabled` default. + */ + getMcpServerOverride(serverId: string): McpServerOverride | undefined { + const override = this.host.activeConversation?.mcpServerOverrides?.find( + (o: McpServerOverride) => o.serverId === serverId + ); + + if (override) return override; + + return this.getDefaultOverride(serverId); + } + + /** + * Gets the effective reasoning effort for the active conversation. + * Returns the conversation override if set, otherwise the global default. + * DEFAULT means no override is sent and the server decides. + */ + getReasoningEffort(): ReasoningEffort { + if (this.host.activeConversation) { + if (this.host.activeConversation.reasoningEffort !== undefined) { + return this.host.activeConversation.reasoningEffort; + } + + // conversations created before the tri-state store an explicit + // opt-out only as thinkingEnabled = false + if (this.host.activeConversation.thinkingEnabled === false) { + return ReasoningEffort.OFF; + } + } + + return this.pendingReasoningEffort; + } + + /** Checks if an MCP server is enabled for the active conversation. */ + isMcpServerEnabledForChat(serverId: string): boolean { + const override = this.getMcpServerOverride(serverId); + + return override?.enabled ?? false; + } + + /** Removes MCP server override for the active conversation. */ + async removeMcpServerOverride(serverId: string): Promise { + await this.setMcpServerOverride(serverId, undefined); + } + + /** Reload persisted defaults, e.g. when the active conversation is cleared. */ + resetPending(): void { + this.pendingReasoningEffort = loadReasoningEffortDefault(); + this.pendingCwd = null; + } + + /** + * Sets the working directory for the active conversation. Pass `null` or + * an empty string to clear it, which restores the picker's empty state. + * + * On the empty new-chat screen (no active conversation yet), the value + * is buffered into `pendingCwd` so the user can pick before + * sending the first message; `createConversation()` consumes it. + * + * @param value - Absolute server-side path to the working directory, or null to clear + */ + async setCwd(value: string | null): Promise { + const trimmed = value?.trim() || undefined; + + // No chat yet - buffer for the first chat the user creates. + if (!this.host.activeConversation) { + this.pendingCwd = trimmed ?? null; + + return; + } + + this.host.applyConversationUpdate(this.host.activeConversation.id, { + cwd: trimmed + }); + + await DatabaseService.updateConversation(this.host.activeConversation.id, { + cwd: trimmed + }); + + this.pendingCwd = null; + } + + /** + * Sets or removes MCP server override for the active conversation. + * If no conversation exists, persists `enabled` onto `mcpServers[i].enabled` + * (the single source of truth for new-chat defaults). + */ + async setMcpServerOverride(serverId: string, enabled: boolean | undefined): Promise { + if (!this.host.activeConversation) { + if (enabled !== undefined) { + mcpStore.updateServer(serverId, { enabled }); + } + + return; + } + + // Clone to plain objects to avoid Proxy serialization issues with IndexedDB + const currentOverrides = (this.host.activeConversation.mcpServerOverrides || []).map( + (o: McpServerOverride) => ({ + enabled: o.enabled, + serverId: o.serverId + }) + ); + + let newOverrides: McpServerOverride[]; + + if (enabled === undefined) { + newOverrides = currentOverrides.filter((o: McpServerOverride) => o.serverId !== serverId); + } else { + const existingIndex = currentOverrides.findIndex( + (o: McpServerOverride) => o.serverId === serverId + ); + + if (existingIndex >= 0) { + newOverrides = [...currentOverrides]; + newOverrides[existingIndex] = { enabled, serverId }; + } else { + newOverrides = [...currentOverrides, { enabled, serverId }]; + } + } + + await DatabaseService.updateConversation(this.host.activeConversation.id, { + mcpServerOverrides: newOverrides.length > 0 ? newOverrides : undefined + }); + + this.host.applyConversationUpdate(this.host.activeConversation.id, { + mcpServerOverrides: newOverrides.length > 0 ? newOverrides : undefined + }); + } + + /** + * Sets the reasoning effort for the active conversation. + * If no conversation exists, stores the global default. + * @param effort - The effort level ('default' | 'off' | 'low' | 'medium' | 'high' | 'max') + */ + async setReasoningEffort(effort: ReasoningEffort): Promise { + if (!this.host.activeConversation) { + this.pendingReasoningEffort = effort; + saveReasoningEffortDefault(effort); + + return; + } + + this.host.applyConversationUpdate(this.host.activeConversation.id, { + reasoningEffort: effort + }); + + await DatabaseService.updateConversation(this.host.activeConversation.id, { + reasoningEffort: effort + }); + } + + /** Toggles MCP server enabled state for the active conversation. */ + async toggleMcpServerForChat(serverId: string): Promise { + const currentEnabled = this.isMcpServerEnabledForChat(serverId); + + await this.setMcpServerOverride(serverId, !currentEnabled); + } + + /** + * Resolve the default enabled value for a server: its own `enabled` + * flag in `mcpServers`, so the global on/off state lives in one place. + */ + private getDefaultOverride(serverId: string): McpServerOverride | undefined { + const server = mcpStore.getServers().find((s) => s.id === serverId); + + if (!server) return undefined; + + return { enabled: server.enabled, serverId }; + } +} diff --git a/tools/ui/src/lib/stores/device.svelte.ts b/tools/ui/src/lib/stores/device.svelte.ts index 08ce2f205428..42aaf4589106 100644 --- a/tools/ui/src/lib/stores/device.svelte.ts +++ b/tools/ui/src/lib/stores/device.svelte.ts @@ -34,11 +34,11 @@ class DeviceStore { readonly isIOSDevice: boolean = false; /** The Safari browser app on iOS, excluding other iOS browsers and WKWebViews. */ readonly isIOSSafari: boolean = false; + /** PWA standalone mode: the page was launched from the home screen icon. */ + isStandalone = $state(false); /** Any WKWebView context on iOS: in-app browsers, embedded web views, and the * third-party iOS browsers (all of which share the WKWebView engine). */ readonly isWKWebView: boolean = false; - /** PWA standalone mode: the page was launched from the home screen icon. */ - isStandalone = $state(false); /** OS color scheme preference; the user override lives in settingsStore. */ readonly systemTheme = $state({ isDark: false }); diff --git a/tools/ui/src/lib/stores/index.ts b/tools/ui/src/lib/stores/index.ts index 1aea8a6dab3e..db227158f45e 100644 --- a/tools/ui/src/lib/stores/index.ts +++ b/tools/ui/src/lib/stores/index.ts @@ -18,34 +18,32 @@ */ // CHAT / MESSAGING -export { chatStore } from './chat.svelte'; +export { chatStore } from './chat/index.svelte'; -export { draftMessagesStore } from './draft-messages.svelte'; +export { draftMessagesStore } from './chat/drafts.svelte'; + +// CONTEXT STATS (active conversation context window usage) +export { contextStatsStore } from './chat/context-stats.svelte'; // AGENTIC (multi-turn tool orchestration) -export { agenticStore } from './agentic.svelte'; +export { agenticStore } from './agentic/index.svelte'; // CONVERSATIONS -export { conversationsStore } from './conversations.svelte'; - -// CONTEXT STATS (active conversation context window usage) -export { contextStatsStore } from './context-stats.svelte'; +export { conversationsStore } from './conversations/index.svelte'; // MCP -export { mcpStore } from './mcp.svelte'; - -export { mcpResourceStore } from './mcp-resources.svelte'; +export { mcpStore } from './mcp/index.svelte'; // MODELS -export { modelsStore } from './models.svelte'; +export { modelsStore } from './models/index.svelte'; // SERVER export { serverStore } from './server.svelte'; // SETTINGS / UI PREFERENCES -export { settingsStore } from './settings.svelte'; +export { settingsStore } from './settings/index.svelte'; -export { settingsReferrer } from './settings-referrer.svelte'; +export { settingsReferrer } from './settings/referrer.svelte'; export { permissionsStore } from './permissions.svelte'; diff --git a/tools/ui/src/lib/stores/init.ts b/tools/ui/src/lib/stores/init.ts index 1faa80303306..d52c34d0feb7 100644 --- a/tools/ui/src/lib/stores/init.ts +++ b/tools/ui/src/lib/stores/init.ts @@ -13,9 +13,9 @@ */ // direct imports, not via the barrel, to avoid circular deps -import { conversationsStore } from './conversations.svelte'; +import { conversationsStore } from './conversations/index.svelte'; import { permissionsStore } from './permissions.svelte'; -import { settingsStore } from './settings.svelte'; +import { settingsStore } from './settings/index.svelte'; import { toolsStore } from './tools.svelte'; import { versionStore } from './version.svelte'; import { browser } from '$app/environment'; @@ -33,7 +33,7 @@ export function initStores(): Promise { permissionsStore.initialize(); toolsStore.initialize(); void versionStore.initialize(); - void conversationsStore.init(); + void conversationsStore.initialize(); })(); return startup; diff --git a/tools/ui/src/lib/stores/mcp/health.svelte.ts b/tools/ui/src/lib/stores/mcp/health.svelte.ts new file mode 100644 index 000000000000..fffa6ea92b00 --- /dev/null +++ b/tools/ui/src/lib/stores/mcp/health.svelte.ts @@ -0,0 +1,298 @@ +/** + * MCPHealthCheckManager - Health checks for MCP servers + * + * Owns per-server connectivity probes: connection reuse, capability + * snapshots, and promotion of a successful check to an active connection. + * Created and owned by mcpStore; the host owns the connection registry the + * probes draw from and promote into. + */ + +import { DEFAULT_MCP_CONFIG } from '$lib/constants'; +import { HealthCheckStatus, MCPConnectionPhase, MCPLogLevel } from '$lib/enums'; +import { MCPService } from '$lib/services/mcp.service'; +import type { + ClientCapabilities, + HealthCheckParams, + HealthCheckState, + MCPCapabilitiesInfo, + MCPConnection, + MCPConnectionLog, + MCPServerConfig, + ServerCapabilities +} from '$lib/types'; +import { detectMcpTransportFromUrl } from '$lib/utils'; + +// module-level so the timestamp is not flagged as reactive state by prefer-svelte-reactivity +function createConnectionErrorLog(message: string): MCPConnectionLog { + return { + level: MCPLogLevel.ERROR, + message: `Connection failed: ${message}`, + phase: MCPConnectionPhase.ERROR, + timestamp: new Date() + }; +} + +/** + * The slice of mcpStore the probes drive. Kept narrow on purpose so the + * probes cannot reach around the host's full surface; mcpStore implements + * this structurally. + */ +export interface McpHealthHost { + autoReconnect(serverName: string): Promise; + getExistingConnection(serverId: string): MCPConnection | undefined; + getRequestTimeoutMs(): number; + promoteHealthCheckToConnection(serverId: string, connection: MCPConnection): void; + registerServerConfig(name: string, config: MCPServerConfig): void; + removeConnection(serverId: string): void; +} + +export class MCPHealthCheckManager { + private _checks = $state>({}); + + /** Raw per-server check states, for host-side capability scans. */ + get checks(): Record { + return this._checks; + } + + clear(serverId: string): void { + const { [serverId]: _removed, ...rest } = this._checks; + + this._checks = rest; + } + + constructor(private host: McpHealthHost) {} + + getState(serverId: string): HealthCheckState { + return this._checks[serverId] ?? { status: HealthCheckStatus.IDLE }; + } + + hasState(serverId: string): boolean { + return serverId in this._checks && this._checks[serverId].status !== HealthCheckStatus.IDLE; + } + + /** + * Run a health check for a server. + * If the server already has an active connection, reuses it instead of creating a new one. + * If promoteToActive is true and server is enabled, the connection will be kept + * and promoted to an active connection instead of being disconnected. + */ + async run(server: HealthCheckParams, promoteToActive = false): Promise { + const existingConnection = this.host.getExistingConnection(server.id); + + if (existingConnection) { + // Reuse existing connection - just refresh tools list + try { + const tools = await MCPService.listTools(existingConnection); + const capabilities = this.buildCapabilitiesInfo( + existingConnection.serverCapabilities, + existingConnection.clientCapabilities + ); + + this.setState(server.id, { + capabilities, + connectionTimeMs: existingConnection.connectionTimeMs, + instructions: existingConnection.instructions, + logs: [], + protocolVersion: existingConnection.protocolVersion, + serverInfo: existingConnection.serverInfo, + status: HealthCheckStatus.SUCCESS, + tools: tools.map((tool) => ({ + description: tool.description, + name: tool.name, + title: tool.title + })), + transportType: existingConnection.transportType + }); + + return; + } catch (error) { + console.warn( + `[MCPStore] Failed to reuse connection for ${server.id}, creating new one:`, + error + ); + // Connection may be stale, remove it and create new one + this.host.removeConnection(server.id); + } + } + + const trimmedUrl = server.url.trim(); + const logs: MCPConnectionLog[] = []; + + let currentPhase: MCPConnectionPhase = MCPConnectionPhase.IDLE; + + if (!trimmedUrl) { + this.setState(server.id, { + logs: [], + message: 'Please enter a server URL first.', + status: HealthCheckStatus.ERROR + }); + + return; + } + + this.setState(server.id, { + logs: [], + phase: MCPConnectionPhase.TRANSPORT_CREATING, + status: HealthCheckStatus.CONNECTING + }); + + const timeoutMs = this.host.getRequestTimeoutMs(); + const headers = this.parseHeaders(server.headers); + + try { + const serverConfig: MCPServerConfig = { + handshakeTimeoutMs: DEFAULT_MCP_CONFIG.connectionTimeoutMs, + headers, + requestTimeoutMs: timeoutMs, + transport: detectMcpTransportFromUrl(trimmedUrl), + url: trimmedUrl, + useProxy: server.useProxy + }; + + this.host.registerServerConfig(server.id, serverConfig); + + const connection = await MCPService.connect( + server.id, + serverConfig, + DEFAULT_MCP_CONFIG.clientInfo, + DEFAULT_MCP_CONFIG.capabilities, + (phase, log) => { + currentPhase = phase; + logs.push(log); + this.setState(server.id, { + logs: [...logs], + phase, + status: HealthCheckStatus.CONNECTING + }); + + if (phase === MCPConnectionPhase.DISCONNECTED && promoteToActive) { + console.log( + `[MCPStore][${server.id}] Connection lost during health check, starting auto-reconnect` + ); + this.host.autoReconnect(server.id); + } + } + ); + const tools = connection.tools.map((tool) => ({ + description: tool.description, + name: tool.name, + title: tool.title + })); + const capabilities = this.buildCapabilitiesInfo( + connection.serverCapabilities, + connection.clientCapabilities + ); + + this.setState(server.id, { + capabilities, + connectionTimeMs: connection.connectionTimeMs, + instructions: connection.instructions, + logs, + protocolVersion: connection.protocolVersion, + serverInfo: connection.serverInfo, + status: HealthCheckStatus.SUCCESS, + tools, + transportType: connection.transportType + }); + + if (promoteToActive && server.enabled) { + this.host.promoteHealthCheckToConnection(server.id, connection); + } else { + await MCPService.disconnect(connection); + } + } catch (error) { + const message = error instanceof Error ? error.message : 'Unknown error occurred'; + + if (logs.at(-1)?.phase !== MCPConnectionPhase.ERROR) { + logs.push(createConnectionErrorLog(message)); + } + + this.setState(server.id, { + logs, + message, + phase: currentPhase, + status: HealthCheckStatus.ERROR + }); + } + } + + async runForServers( + servers: { + id: string; + enabled: boolean; + url: string; + headers?: string; + }[], + skipIfChecked = true, + promoteToActive = false + ): Promise { + const serversToCheck = skipIfChecked + ? servers.filter((s) => !this.hasState(s.id) && s.url.trim()) + : servers.filter((s) => s.url.trim()); + + if (serversToCheck.length === 0) { + return; + } + + const BATCH_SIZE = 5; + + for (let i = 0; i < serversToCheck.length; i += BATCH_SIZE) { + const batch = serversToCheck.slice(i, i + BATCH_SIZE); + + await Promise.allSettled(batch.map((server) => this.run(server, promoteToActive))); + } + } + + /** + * Builds capabilities info from server and client capabilities. + */ + private buildCapabilitiesInfo( + serverCaps?: ServerCapabilities, + clientCaps?: ClientCapabilities + ): MCPCapabilitiesInfo { + return { + client: { + elicitation: clientCaps?.elicitation + ? { form: !!clientCaps.elicitation.form, url: !!clientCaps.elicitation.url } + : undefined, + roots: clientCaps?.roots ? { listChanged: clientCaps.roots.listChanged } : undefined, + sampling: !!clientCaps?.sampling, + tasks: !!clientCaps?.tasks + }, + server: { + completions: !!serverCaps?.completions, + logging: !!serverCaps?.logging, + prompts: serverCaps?.prompts ? { listChanged: serverCaps.prompts.listChanged } : undefined, + resources: serverCaps?.resources + ? { + listChanged: serverCaps.resources.listChanged, + subscribe: serverCaps.resources.subscribe + } + : undefined, + tasks: !!serverCaps?.tasks, + tools: serverCaps?.tools ? { listChanged: serverCaps.tools.listChanged } : undefined + } + }; + } + + private parseHeaders(headersJson?: string): Record | undefined { + if (!headersJson?.trim()) { + return undefined; + } + + try { + const parsed = JSON.parse(headersJson); + + if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) + return parsed as Record; + } catch { + console.warn('[MCPStore] Failed to parse custom headers JSON:', headersJson); + } + + return undefined; + } + + private setState(serverId: string, state: HealthCheckState): void { + this._checks = { ...this._checks, [serverId]: state }; + } +} diff --git a/tools/ui/src/lib/stores/mcp.svelte.ts b/tools/ui/src/lib/stores/mcp/index.svelte.ts similarity index 67% rename from tools/ui/src/lib/stores/mcp.svelte.ts rename to tools/ui/src/lib/stores/mcp/index.svelte.ts index 3e0cb8e1e39e..ccd53bc9d2f9 100644 --- a/tools/ui/src/lib/stores/mcp.svelte.ts +++ b/tools/ui/src/lib/stores/mcp/index.svelte.ts @@ -1,69 +1,38 @@ /** - * mcpStore - Reactive State Store for MCP Operations + * mcpStore - MCP host: server connections and tool operations * - * Implements the "Host" role in MCP architecture, coordinating multiple server - * connections and providing a unified interface for tool operations. - * - * **Architecture & Relationships:** - * - **MCPService**: Stateless protocol layer (transport, connect, callTool) - * - **mcpStore** (this): Reactive state + business logic - * - * **Key Responsibilities:** - * - Lifecycle management (initialize, shutdown) - * - Multi-server coordination - * - Tool name conflict detection and resolution - * - Automatic tool-to-server routing - * - Health checks - * - * MCP connection state and raw `Tool[]` per server are owned here; the - * OpenAI-compatible wire format for those tools is built in `toolsStore` - * (see {@link toolsStore.mcpEntries} / {@link toolsStore.getEnabledToolsForLLM}). - * - * @see MCPService in services/mcp.service.ts for protocol operations + * Implements the MCP "Host" role, coordinating multiple server connections + * and exposing a unified tool interface: lifecycle, name-conflict detection + * and automatic tool-to-server routing. Owns connection state and raw + * `Tool[]` per server; the OpenAI-compatible wire format is built in + * toolsStore. Composes the health-check manager; uses MCPService for the + * protocol layer. */ import type { ListChangedHandlers } from '@modelcontextprotocol/sdk/types.js'; import { browser } from '$app/environment'; import { SETTINGS_KEYS } from '$lib/constants'; -import { - CACHE, - DEFAULT_MCP_CONFIG, - EXPECTED_THEMED_ICON_PAIR_COUNT, - MCP_ALLOWED_ICON_MIME_TYPES, - MCP_RECONNECT, - MCP_SERVER_ID_PREFIX -} from '$lib/constants'; -import { - ColorMode, - HealthCheckStatus, - MCPConnectionPhase, - MCPLogLevel, - MCPRefType, - UrlProtocol -} from '$lib/enums'; +import { CACHE, DEFAULT_MCP_CONFIG, MCP_RECONNECT, MCP_SERVER_ID_PREFIX } from '$lib/constants'; +import { ColorMode, HealthCheckStatus, MCPConnectionPhase, MCPRefType } from '$lib/enums'; import { MCPService } from '$lib/services/mcp.service'; // direct imports between stores, not via the barrel, to avoid circular deps -import { mcpResourceStore } from '$lib/stores/mcp-resources.svelte'; +import { MCPHealthCheckManager, type McpHealthHost } from '$lib/stores/mcp/health.svelte'; +import { mcpResourceStore } from '$lib/stores/mcp/resources.svelte'; import { serverStore } from '$lib/stores/server.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { - ClientCapabilities, GetPromptResult, HealthCheckParams, HealthCheckState, - MCPCapabilitiesInfo, MCPClientConfig, MCPConnection, - MCPConnectionLog, MCPPromptInfo, MCPResourceAttachment, MCPResourceContent, - MCPResourceIcon, MCPServerConfig, MCPServerDisplayInfo, MCPServerSettingsEntry, MCPToolCall, - ServerCapabilities, ServerStatus, Tool, ToolExecutionResult @@ -72,541 +41,269 @@ import type { DatabaseMessageExtraMcpResource, McpServerOverride } from '$lib/ty import type { SettingsConfigType } from '$lib/types/settings'; import { detectMcpTransportFromUrl, - extractRootDomain, + getMcpIconUrl, + getMcpServerFaviconFallback, + getMcpServerLabel, parseMcpServerSettings, uuid } from '$lib/utils'; import { mode } from 'mode-watcher'; -class MCPStore { - private _isInitializing = $state(false); +class MCPStore implements McpHealthHost { private _error = $state(null); + private _isInitializing = $state(false); private _toolCount = $state(0); - private _connectedServers = $state([]); - private _healthChecks = $state>({}); + private activeFlowCount = 0; - private connections = new Map(); - private toolsIndex = new Map(); - private serverConfigs = new Map(); // Store configs for reconnection - private reconnectingServers = new Set(); // Guard against concurrent reconnections private configSignature: string | null = null; + private connectedServers = $state([]); + private connections = new Map(); + // health checks: per-server connectivity probes with optional promotion to active connections + private health = new MCPHealthCheckManager(this); private initPromise: Promise | null = null; - private activeFlowCount = 0; - - get isProxyAvailable(): boolean { - return serverStore.props?.cors_proxy_enabled ?? false; - } - - /** - * Generates a unique server ID from an optional ID string or index. - */ - #generateServerId(id: unknown, index: number): string { - if (typeof id === 'string' && id.trim()) { - return id.trim(); - } - - return `${MCP_SERVER_ID_PREFIX}-${index + 1}`; - } - - /** - * Parses raw server settings from config into MCPServerSettingsEntry array. - */ - #parseServerSettings(rawServers: unknown): MCPServerSettingsEntry[] { - if (!rawServers) { - return []; - } - - let parsed: unknown; - - if (typeof rawServers === 'string') { - const trimmed = rawServers.trim(); - - if (!trimmed) { - return []; - } - - try { - parsed = JSON.parse(trimmed); - } catch (error) { - console.warn('[MCP] Failed to parse mcpServers JSON:', error); - - return []; - } - } else { - parsed = rawServers; - } - - if (!Array.isArray(parsed)) { - return []; - } - - return parsed.map((entry, index) => { - const url = typeof entry?.url === 'string' ? entry.url.trim() : ''; - const headers = typeof entry?.headers === 'string' ? entry.headers.trim() : undefined; + private reconnectingServers = new Set(); // Guard against concurrent reconnections + private serverConfigs = new Map(); // Store configs for reconnection + private serversCache: { raw: unknown; servers: MCPServerSettingsEntry[] } | null = null; + private toolsIndex = new Map(); - return { - displayName: (entry as { displayName?: string })?.displayName, - enabled: Boolean((entry as { enabled?: unknown })?.enabled), - headers: headers || undefined, - id: this.#generateServerId((entry as { id?: unknown })?.id, index), - name: (entry as { name?: string })?.name, - url, - useProxy: Boolean((entry as { useProxy?: unknown })?.useProxy) - } satisfies MCPServerSettingsEntry; - }); + get availableTools(): string[] { + return Array.from(this.toolsIndex.keys()); } - /** - * Request timeout in milliseconds, read live from the global setting - * so a change in Settings applies to every server immediately. - */ - #requestTimeoutMs(): number { - const seconds = - Number(settingsStore.config.mcpRequestTimeoutSeconds) || - DEFAULT_MCP_CONFIG.requestTimeoutSeconds; - - return Math.round(seconds * 1000); + get connectedServerCount(): number { + return this.connectedServers.length; } - /** - * Builds server configuration from a settings entry. - */ - #buildServerConfig( - entry: MCPServerSettingsEntry, - connectionTimeoutMs = DEFAULT_MCP_CONFIG.connectionTimeoutMs - ): MCPServerConfig | undefined { - if (!entry?.url) { - return undefined; - } - - let headers: Record | undefined; - - if (entry.headers) { - try { - const parsed = JSON.parse(entry.headers); - - if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) - headers = parsed as Record; - } catch { - console.warn('[MCP] Failed to parse custom headers JSON:', entry.headers); - } - } - - return { - handshakeTimeoutMs: connectionTimeoutMs, - headers, - requestTimeoutMs: this.#requestTimeoutMs(), - transport: detectMcpTransportFromUrl(entry.url), - url: entry.url, - useProxy: entry.useProxy - }; + get connectedServerNames(): string[] { + return this.connectedServers; } - /** - * Checks if a server is enabled for a given chat. - * A per-chat override wins when present; a server without one resolves - * to its own `enabled` flag in `mcpServers`. - */ - #checkServerEnabled( - server: MCPServerSettingsEntry, - perChatOverrides?: McpServerOverride[] - ): boolean { - // Per-chat overrides win when present; missing entries inherit the - // server's own `enabled` flag so partial override lists are not all - // treated as disabled. - const override = perChatOverrides?.find((o) => o.serverId === server.id); - - return override?.enabled ?? server.enabled; + get error(): string | null { + return this._error; } - /** - * Builds MCP client configuration from settings. - */ - #buildMcpClientConfig( - cfg: SettingsConfigType, - perChatOverrides?: McpServerOverride[] - ): MCPClientConfig | undefined { - const rawServers = this.#parseServerSettings(cfg.mcpServers); - - if (!rawServers.length) { - return undefined; - } - - const servers: Record = {}; - - for (const [index, entry] of rawServers.entries()) { - if (!this.#checkServerEnabled(entry, perChatOverrides)) continue; - - const normalized = this.#buildServerConfig(entry); - - if (normalized) servers[this.#generateServerId(entry.id, index)] = normalized; - } - - if (Object.keys(servers).length === 0) { - return undefined; - } + get isEnabled(): boolean { + const mcpConfig = this.buildMcpClientConfig(settingsStore.config); - return { - capabilities: DEFAULT_MCP_CONFIG.capabilities, - clientInfo: DEFAULT_MCP_CONFIG.clientInfo, - protocolVersion: DEFAULT_MCP_CONFIG.protocolVersion, - requestTimeoutMs: this.#requestTimeoutMs(), - servers - }; + return ( + mcpConfig !== null && mcpConfig !== undefined && Object.keys(mcpConfig.servers).length > 0 + ); } - /** - * Builds capabilities info from server and client capabilities. - */ - #buildCapabilitiesInfo( - serverCaps?: ServerCapabilities, - clientCaps?: ClientCapabilities - ): MCPCapabilitiesInfo { - return { - client: { - elicitation: clientCaps?.elicitation - ? { form: !!clientCaps.elicitation.form, url: !!clientCaps.elicitation.url } - : undefined, - roots: clientCaps?.roots ? { listChanged: clientCaps.roots.listChanged } : undefined, - sampling: !!clientCaps?.sampling, - tasks: !!clientCaps?.tasks - }, - server: { - completions: !!serverCaps?.completions, - logging: !!serverCaps?.logging, - prompts: serverCaps?.prompts ? { listChanged: serverCaps.prompts.listChanged } : undefined, - resources: serverCaps?.resources - ? { - listChanged: serverCaps.resources.listChanged, - subscribe: serverCaps.resources.subscribe - } - : undefined, - tasks: !!serverCaps?.tasks, - tools: serverCaps?.tools ? { listChanged: serverCaps.tools.listChanged } : undefined - } - }; + get isInitialized(): boolean { + return this.connections.size > 0; } get isInitializing(): boolean { return this._isInitializing; } - get isInitialized(): boolean { - return this.connections.size > 0; + get isProxyAvailable(): boolean { + return serverStore.props?.cors_proxy_enabled ?? false; } - get error(): string | null { - return this._error; + /** Resource state, composed here so consumers have a single MCP scope. */ + get resources() { + return mcpResourceStore; } get toolCount(): number { return this._toolCount; } - get connectedServerCount(): number { - return this._connectedServers.length; - } - - get connectedServerNames(): string[] { - return this._connectedServers; + acquireConnection(): void { + this.activeFlowCount++; } - get isEnabled(): boolean { - const mcpConfig = this.#buildMcpClientConfig(settingsStore.config); + addServer( + serverData: Omit & { id?: string } + ): MCPServerSettingsEntry { + const servers = this.getServers(); + const newServer: MCPServerSettingsEntry = { + displayName: serverData.displayName, + enabled: serverData.enabled, + headers: serverData.headers?.trim() || undefined, + id: serverData.id || (uuid() ?? `server-${Date.now()}`), + name: serverData.name, + url: serverData.url.trim(), + useProxy: serverData.useProxy + }; - return ( - mcpConfig !== null && mcpConfig !== undefined && Object.keys(mcpConfig.servers).length > 0 - ); - } + settingsStore.updateConfig(SETTINGS_KEYS.MCP_SERVERS, JSON.stringify([...servers, newServer])); - get availableTools(): string[] { - return Array.from(this.toolsIndex.keys()); + return newServer; } - private updateState(state: { - isInitializing?: boolean; - error?: string | null; - toolCount?: number; - connectedServers?: string[]; - }): void { - if (state.isInitializing !== undefined) { - this._isInitializing = state.isInitializing; - } + /** + * Add a resource as attachment to chat context. + * Automatically fetches content if not cached. + */ + async attachResource(uri: string): Promise { + const resourceInfo = mcpResourceStore.findResourceByUri(uri); - if (state.error !== undefined) { - this._error = state.error; - } + if (!resourceInfo) { + console.error(`[MCPStore] Resource not found: ${uri}`); - if (state.toolCount !== undefined) { - this._toolCount = state.toolCount; + return null; } - if (state.connectedServers !== undefined) { - this._connectedServers = state.connectedServers; + if (mcpResourceStore.isAttached(uri)) { + return null; } - } - updateHealthCheck(serverId: string, state: HealthCheckState): void { - this._healthChecks = { ...this._healthChecks, [serverId]: state }; - } - - getHealthCheckState(serverId: string): HealthCheckState { - return this._healthChecks[serverId] ?? { status: HealthCheckStatus.IDLE }; - } - - hasHealthCheck(serverId: string): boolean { - return ( - serverId in this._healthChecks && - this._healthChecks[serverId].status !== HealthCheckStatus.IDLE - ); - } - - clearHealthCheck(serverId: string): void { - const { [serverId]: _removed, ...rest } = this._healthChecks; + const attachment = mcpResourceStore.addAttachment(resourceInfo); - this._healthChecks = rest; - } + try { + const content = await this.readResource(uri); - clearAllHealthChecks(): void { - this._healthChecks = {}; - } + if (content) { + mcpResourceStore.updateAttachmentContent(attachment.id, content); + } else { + mcpResourceStore.updateAttachmentError(attachment.id, 'Failed to read resource'); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); - clearError(): void { - this._error = null; - } + mcpResourceStore.updateAttachmentError(attachment.id, message); + } - getServers(): MCPServerSettingsEntry[] { - return parseMcpServerSettings(settingsStore.config.mcpServers); + return mcpResourceStore.getAttachment(attachment.id) ?? null; } /** - * Get all active MCP connections. - * @returns Map of server names to connections + * Auto-reconnect to a server with exponential backoff. + * Continues indefinitely until successful. + * + * Race-condition safety: when the phase callback fires a DISCONNECTED event + * while we are still inside this function (e.g., the server drops right after + * a successful connect()), a naive inner `autoReconnect()` call would be + * swallowed by the `reconnectingServers` guard, leaving the server + * permanently disconnected once the outer call exits. We solve this by + * deferring the new reconnection via the `needsReconnect` flag: the flag is + * set inside the phase callback and honoured in the `finally` block after + * the guard entry has been removed. */ - getConnections(): Map { - return this.connections; - } + async autoReconnect(serverName: string): Promise { + // Guard against concurrent reconnections + if (this.reconnectingServers.has(serverName)) { + console.log(`[MCPStore][${serverName}] Reconnection already in progress, skipping`); - /** - * Resolves the raw label for a server: user-defined display name first, - * then server-reported title or name when the health check succeeded, - * then the configured name (admin baseline or legacy data), then URL. - */ - #serverBaseLabel(server: MCPServerDisplayInfo): string { - if (server.displayName) return server.displayName; - - const healthState = this.getHealthCheckState(server.id); - - if (healthState?.status === HealthCheckStatus.SUCCESS) - return ( - healthState.serverInfo?.title || healthState.serverInfo?.name || server.name || server.url - ); - - return server.name || server.url; - } - - /** - * Returns the display label for a server, suffixed with a positional - * counter when several configured servers resolve to the same base label - * (e.g. two endpoints of the same host reporting an identical name). - * Numbering follows config order, so it is stable across renders. - */ - getServerLabel(server: MCPServerDisplayInfo): string { - const label = this.#serverBaseLabel(server); - const twins = this.getServers().filter((s) => this.#serverBaseLabel(s) === label); - - if (twins.length < 2) return label; - - const position = twins.findIndex((s) => s.id === server.id); + return; + } - return position < 0 ? label : `${label} (${position + 1})`; - } + const serverConfig = this.serverConfigs.get(serverName); - getServerById(serverId: string): MCPServerSettingsEntry | undefined { - return this.getServers().find((s) => s.id === serverId); - } + if (!serverConfig) { + console.error(`[MCPStore] No config found for ${serverName}, cannot reconnect`); - /** - * Get display name for an MCP server by its ID. - * Falls back to the server ID if server is not found. - */ - getServerDisplayName(serverId: string): string { - const server = this.getServerById(serverId); + return; + } - return server ? this.getServerLabel(server) : serverId; - } + this.reconnectingServers.add(serverName); + let backoff = MCP_RECONNECT.INITIAL_DELAY; + // Flag set by the phase callback when a DISCONNECTED event fires while + // reconnectingServers still holds this server (see JSDoc above). + let needsReconnect = false; - /** - * Validates that an icon URI uses a safe scheme (https: or data:). - */ - #isValidIconUri(src: string): boolean { try { - if (src.startsWith(UrlProtocol.DATA)) return true; - - const url = new URL(src); - - return url.protocol === UrlProtocol.HTTPS; - } catch { - return false; - } - } - - /** - * Selects the best icon URL from an MCP icons array. - * Follows security guidelines from the MCP specification: - * - Only allows https: and data: URIs - * - Filters to supported MIME types - * - * Selection priority: - * 1. Icon matching the current color scheme (dark/light) - * 2. Universal icon (no theme specified); if exactly 2, assumes [0]=light, [1]=dark - * 3. First valid icon as last resort - */ - #getMcpIconUrl(icons: MCPResourceIcon[] | undefined, isDark = false): string | null { - if (!icons?.length) return null; + while (true) { + await new Promise((resolve) => setTimeout(resolve, backoff)); - const validIcons = icons.filter((icon) => { - if (!icon.src || !this.#isValidIconUri(icon.src)) return false; + console.log(`[MCPStore][${serverName}] Auto-reconnecting...`); - if (icon.mimeType && !MCP_ALLOWED_ICON_MIME_TYPES.has(icon.mimeType)) return false; + try { + // Per-attempt timeout: reject if the server doesn't respond in time, + // then fall through to backoff logic as with any other failure. + const timeoutPromise = new Promise((_, reject) => + setTimeout( + () => + reject( + new Error( + `Reconnect attempt timed out after ${MCP_RECONNECT.ATTEMPT_TIMEOUT_MS}ms` + ) + ), + MCP_RECONNECT.ATTEMPT_TIMEOUT_MS + ) + ); - return true; - }); + needsReconnect = false; + const listChangedHandlers = this.createListChangedHandlers(serverName); + const connectPromise = MCPService.connect( + serverName, + serverConfig, + DEFAULT_MCP_CONFIG.clientInfo, + DEFAULT_MCP_CONFIG.capabilities, + (phase) => { + if (phase === MCPConnectionPhase.DISCONNECTED) { + if (this.reconnectingServers.has(serverName)) { + // Reconnect loop is active; defer to after it exits. + needsReconnect = true; + } else { + console.log( + `[MCPStore][${serverName}] Connection lost, restarting auto-reconnect` + ); + this.autoReconnect(serverName); + } + } + }, + listChangedHandlers + ); + const connection = await Promise.race([connectPromise, timeoutPromise]); - if (validIcons.length === 0) return null; + this.connections.set(serverName, connection); - const preferredTheme = isDark ? ColorMode.DARK : ColorMode.LIGHT; - // 1. Prefer icon explicitly matching the current color scheme - const themedIcon = validIcons.find((icon) => icon.theme === preferredTheme); + // Rebuild tool index for this server + this.indexServerTools(serverName, connection.tools); - if (themedIcon) return themedIcon.src; + console.log(`[MCPStore][${serverName}] Reconnected successfully`); - // 2. Handle universal icons (no theme specified) - const universalIcons = validIcons.filter((icon) => !icon.theme); + break; + } catch (error) { + console.warn(`[MCPStore][${serverName}] Reconnection failed:`, error); + backoff = Math.min(backoff * MCP_RECONNECT.BACKOFF_MULTIPLIER, MCP_RECONNECT.MAX_DELAY); + } + } + } finally { + this.reconnectingServers.delete(serverName); - if (universalIcons.length === EXPECTED_THEMED_ICON_PAIR_COUNT) { - // Heuristic: two theme-less icons → assume [0] = light, [1] = dark - return universalIcons[isDark ? 1 : 0].src; + // If the phase callback signalled a disconnect while this function held + // the guard, kick off a fresh reconnect now that the guard is released. + if (needsReconnect) { + console.log( + `[MCPStore][${serverName}] Deferred disconnect detected, restarting auto-reconnect` + ); + this.autoReconnect(serverName); + } } + } - if (universalIcons.length > 0) { - return universalIcons[0].src; - } + clearError(): void { + this._error = null; + } - // 3. Last resort: use opposite-theme icon - return validIcons[0].src; + clearHealthCheck(serverId: string): void { + this.health.clear(serverId); } /** - * Get icon URL for an MCP server by its ID. - * Returns the best icon from the MCP server's `icons` array - * (see MCP spec: spec.modelcontextprotocol.io). - * Returns null if no icon is available. + * Clear all resource attachments. */ - getServerFavicon(serverId: string): string | null { - const server = this.getServerById(serverId); - - if (!server) { - return null; - } - - const isDark = mode.current === ColorMode.DARK; - const healthState = this.getHealthCheckState(serverId); - - if (healthState.status === HealthCheckStatus.SUCCESS && healthState.serverInfo?.icons) { - const mcpIconUrl = this.#getMcpIconUrl(healthState.serverInfo.icons, isDark); - - if (mcpIconUrl) { - return mcpIconUrl; - } - } - - return this.#getServerFaviconFallback(server.url); + clearResourceAttachments(): void { + mcpResourceStore.clearAttachments(); } /** - * Construct a fallback favicon URL from the MCP server URL. - * e.g. https://mcp.example.com/sse -> https://example.com/favicon.ico + * Convert current resource attachments to DatabaseMessageExtra[] and clear them. + * Called during message send to persist resources with the user message. */ - #getServerFaviconFallback(serverUrl: string): string | null { - try { - const url = new URL(serverUrl); - const rootDomain = extractRootDomain(url); - - if (!rootDomain) return null; - - const origin = `${url.protocol}//${rootDomain}`; - const candidates = ['favicon.ico', 'favicon.png']; - - for (const path of candidates) { - const faviconUrl = `${origin}/${path}`; + consumeResourceAttachmentsAsExtras(): DatabaseMessageExtraMcpResource[] { + const extras = mcpResourceStore.toMessageExtras(); - if (this.#isValidIconUri(faviconUrl)) { - return faviconUrl; - } - } - } catch { - // Invalid URL, return null + if (extras.length > 0) { + mcpResourceStore.clearAttachments(); } - return null; - } - - addServer( - serverData: Omit & { id?: string } - ): MCPServerSettingsEntry { - const servers = this.getServers(); - const newServer: MCPServerSettingsEntry = { - displayName: serverData.displayName, - enabled: serverData.enabled, - headers: serverData.headers?.trim() || undefined, - id: serverData.id || (uuid() ?? `server-${Date.now()}`), - name: serverData.name, - url: serverData.url.trim(), - useProxy: serverData.useProxy - }; - - settingsStore.updateConfig(SETTINGS_KEYS.MCP_SERVERS, JSON.stringify([...servers, newServer])); - - return newServer; - } - - updateServer(id: string, updates: Partial): void { - const servers = this.getServers(); - - settingsStore.updateConfig( - SETTINGS_KEYS.MCP_SERVERS, - JSON.stringify( - servers.map((server) => (server.id === id ? { ...server, ...updates } : server)) - ) - ); - } - - removeServer(id: string): void { - const servers = this.getServers(); - - settingsStore.updateConfig( - SETTINGS_KEYS.MCP_SERVERS, - JSON.stringify(servers.filter((s) => s.id !== id)) - ); - this.clearHealthCheck(id); - } - - hasAvailableServers(): boolean { - return parseMcpServerSettings(settingsStore.config.mcpServers).some( - (s) => s.enabled && s.url.trim() - ); - } - hasEnabledServers(perChatOverrides?: McpServerOverride[]): boolean { - return Boolean(this.#buildMcpClientConfig(settingsStore.config, perChatOverrides)); - } - - getEnabledServersForConversation( - perChatOverrides?: McpServerOverride[] - ): MCPServerSettingsEntry[] { - return this.getServers().filter((server) => { - return this.#checkServerEnabled(server, perChatOverrides); - }); + return extras; } async ensureInitialized(perChatOverrides?: McpServerOverride[]): Promise { @@ -614,7 +311,7 @@ class MCPStore { return false; } - const mcpConfig = this.#buildMcpClientConfig(settingsStore.config, perChatOverrides); + const mcpConfig = this.buildMcpClientConfig(settingsStore.config, perChatOverrides); const signature = mcpConfig ? JSON.stringify(mcpConfig) : null; if (!signature) { @@ -636,403 +333,357 @@ class MCPStore { return this.initialize(signature, mcpConfig!); } - private async initialize(signature: string, mcpConfig: MCPClientConfig): Promise { - this.updateState({ error: null, isInitializing: true }); - this.configSignature = signature; + async executeTool(toolCall: MCPToolCall, signal?: AbortSignal): Promise { + return this.executeToolByName( + toolCall.function.name, + this.parseToolArguments(toolCall.function.arguments), + signal + ); + } - const serverEntries = Object.entries(mcpConfig.servers); + async executeToolByName( + toolName: string, + args: Record, + signal?: AbortSignal + ): Promise { + const serverName = this.toolsIndex.get(toolName); - if (serverEntries.length === 0) { - this.updateState({ connectedServers: [], isInitializing: false, toolCount: 0 }); + if (!serverName) throw new Error(`Unknown tool: ${toolName}`); - return false; - } + const connection = this.connections.get(serverName); - this.initPromise = this.doInitialize(signature, mcpConfig, serverEntries); + if (!connection) throw new Error(`Server "${serverName}" is not connected`); - return this.initPromise; - } + try { + return await MCPService.callTool(connection, { arguments: args, name: toolName }, signal); + } catch (error) { + if (MCPService.isSessionExpiredError(error)) { + await this.reconnectServer(serverName); - private async doInitialize( - signature: string, - mcpConfig: MCPClientConfig, - serverEntries: [string, MCPClientConfig['servers'][string]][] - ): Promise { - const clientInfo = mcpConfig.clientInfo ?? DEFAULT_MCP_CONFIG.clientInfo; - const capabilities = mcpConfig.capabilities ?? DEFAULT_MCP_CONFIG.capabilities; - const results = await Promise.allSettled( - serverEntries.map(async ([name, serverConfig]) => { - // Store config for reconnection - this.serverConfigs.set(name, serverConfig); + const newConnection = this.connections.get(serverName); - const listChangedHandlers = this.createListChangedHandlers(name); - const connection = await MCPService.connect( - name, - serverConfig, - clientInfo, - capabilities, - (phase) => { - // Handle WebSocket disconnection - if (phase === MCPConnectionPhase.DISCONNECTED) { - console.log(`[MCPStore][${name}] Connection lost, starting auto-reconnect`); - this.autoReconnect(name); - } - }, - listChangedHandlers - ); - - return { connection, name }; - }) - ); + if (!newConnection) throw new Error(`Failed to reconnect to "${serverName}"`); - if (this.configSignature !== signature) { - for (const result of results) { - if (result.status === 'fulfilled') - await MCPService.disconnect(result.value.connection).catch(console.warn); + return MCPService.callTool(newConnection, { arguments: args, name: toolName }, signal); } - return false; + throw error; } + } - for (const result of results) { - if (result.status === 'fulfilled') { - const { connection, name } = result.value; + /** + * Fetch resources from all connected servers that support them. + * Updates mcpResourceStore with the results. + * @param forceRefresh - If true, bypass cache and fetch fresh data + */ + async fetchAllResources(forceRefresh: boolean = false): Promise { + const serversWithResources = this.getServersWithResources(); - this.connections.set(name, connection); + if (serversWithResources.length === 0) { + return; + } - for (const tool of connection.tools) { - if (this.toolsIndex.has(tool.name)) - console.warn( - `[MCPStore] Tool name conflict: "${tool.name}" exists in "${this.toolsIndex.get(tool.name)}" and "${name}". Using tool from "${name}".` - ); + // Check if we have cached resources and they're recent (unless force refresh) + if (!forceRefresh) { + const allServersCached = serversWithResources.every((serverName) => { + const serverRes = mcpResourceStore.getServerResources(serverName); - this.toolsIndex.set(tool.name, name); + if (!serverRes || !serverRes.lastFetched) { + return false; } - } else { - console.error(`[MCPStore] Failed to connect:`, result.reason); - } - } - const successCount = this.connections.size; + // Cache is valid for 5 minutes + const age = Date.now() - serverRes.lastFetched.getTime(); - if (successCount === 0 && serverEntries.length > 0) { - this.updateState({ - connectedServers: [], - error: 'All MCP server connections failed', - isInitializing: false, - toolCount: 0 + return age < CACHE.DEFAULT_TTL_MS; }); - this.initPromise = null; - - return false; - } - - this.updateState({ - connectedServers: Array.from(this.connections.keys()), - error: null, - isInitializing: false, - toolCount: this.toolsIndex.size - }); - this.initPromise = null; - - return true; - } - private createListChangedHandlers(serverName: string): ListChangedHandlers { - return { - prompts: { - onChanged: (error: Error | null) => { - if (error) { - console.warn(`[MCPStore][${serverName}] Prompts list changed error:`, error); + if (allServersCached) { + console.log('[MCPStore] Using cached resources'); - return; - } - } - }, - tools: { - onChanged: (error: Error | null, tools: Tool[] | null) => { - if (error) { - console.warn(`[MCPStore][${serverName}] Tools list changed error:`, error); + return; + } + } - return; - } + mcpResourceStore.setLoading(true); - this.handleToolsListChanged(serverName, tools ?? []); - } - } - }; + try { + await Promise.all( + serversWithResources.map((serverName) => this.fetchServerResources(serverName)) + ); + } finally { + mcpResourceStore.setLoading(false); + } } - private handleToolsListChanged(serverName: string, tools: Tool[]): void { + /** + * Fetch resources from a specific server. + * Updates mcpResourceStore with the results. + */ + async fetchServerResources(serverName: string): Promise { const connection = this.connections.get(serverName); if (!connection) { + console.warn(`[MCPStore] No connection found for server: ${serverName}`); + return; } - for (const [toolName, ownerServer] of this.toolsIndex.entries()) { - if (ownerServer === serverName) this.toolsIndex.delete(toolName); + if (!MCPService.supportsResources(connection)) { + return; } - connection.tools = tools; + mcpResourceStore.setServerLoading(serverName, true); - for (const tool of tools) { - if (this.toolsIndex.has(tool.name)) - console.warn( - `[MCPStore] Tool name conflict after list change: "${tool.name}" exists in "${this.toolsIndex.get(tool.name)}" and "${serverName}". Using tool from "${serverName}".` - ); + try { + const [resources, templates] = await Promise.all([ + MCPService.listAllResources(connection), + MCPService.listAllResourceTemplates(connection) + ]); - this.toolsIndex.set(tool.name, serverName); - } - this.updateState({ toolCount: this.toolsIndex.size }); - } + mcpResourceStore.setServerResources(serverName, resources, templates); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); - acquireConnection(): void { - this.activeFlowCount++; + mcpResourceStore.setServerError(serverName, message); + console.error(`[MCPStore][${serverName}] Failed to fetch resources:`, error); + } } /** - * Release a connection reference. - * By default, keeps connections alive for reuse (shutdownIfUnused=false). - * MCP spec encourages long-lived sessions to avoid reconnection overhead. + * Resolve which configured MCP server owns a given tool name. Looks at + * active connections first (fast path), then falls back to per-server + * health-check data so server-side MCP proxies (where llama-server + * executes MCP tools but the browser does not hold a direct connection) + * still resolve tool names to their owning server. */ - async releaseConnection(shutdownIfUnused = false): Promise { - this.activeFlowCount = Math.max(0, this.activeFlowCount - 1); + findServerForTool(toolName: string): string | undefined { + const fromIndex = this.toolsIndex.get(toolName); - if (shutdownIfUnused && this.activeFlowCount === 0) { - await this.shutdown(); + if (fromIndex) return fromIndex; + + for (const server of this.getServers()) { + const health = this.health.checks[server.id]; + + if (!health || health.status !== HealthCheckStatus.SUCCESS) continue; + + if (health.tools.some((tool) => tool.name === toolName)) { + return server.id; + } } - } + return undefined; + } getActiveFlowCount(): number { return this.activeFlowCount; } - async shutdown(): Promise { - if (this.initPromise) { - await this.initPromise.catch(() => {}); - this.initPromise = null; - } + async getAllPrompts(): Promise { + const results: MCPPromptInfo[] = []; - if (this.connections.size === 0) { - return; + for (const [serverName, connection] of this.connections) { + if (!connection.serverCapabilities?.prompts) continue; + + const prompts = await MCPService.listPrompts(connection); + + for (const prompt of prompts) { + results.push({ + arguments: prompt.arguments?.map((arg) => ({ + description: arg.description, + name: arg.name, + required: arg.required + })), + description: prompt.description, + name: prompt.name, + serverName, + title: prompt.title + }); + } } - await Promise.all( - Array.from(this.connections.values()).map((conn) => - MCPService.disconnect(conn).catch((error) => - console.warn(`[MCPStore] Error disconnecting ${conn.serverName}:`, error) - ) - ) - ); + return results; + } - this.connections.clear(); - this.toolsIndex.clear(); - this.serverConfigs.clear(); - this.configSignature = null; - this.updateState({ - connectedServers: [], - error: null, - isInitializing: false, - toolCount: 0 + /** + * Get all active MCP connections. + * @returns Map of server names to connections + */ + getConnections(): Map { + return this.connections; + } + + getEnabledServersForConversation( + perChatOverrides?: McpServerOverride[] + ): MCPServerSettingsEntry[] { + return this.getServers().filter((server) => { + return this.checkServerEnabled(server, perChatOverrides); }); } /** - * Immediately reconnect to a server by creating a fresh transport and session. - * Used when a session-expired error (HTTP 404) is detected during tool execution. - * Per MCP spec 2025-11-25: client MUST discard session ID and re-initialize. - * - * Unlike autoReconnect (which uses exponential backoff for connectivity issues), - * this performs a single immediate reconnection attempt since the server is known - * to be reachable (it responded with 404). + * Check if a server already has an active connection that can be reused. + * Returns the existing connection if available. */ - private async reconnectServer(serverName: string): Promise { - const serverConfig = this.serverConfigs.get(serverName); + getExistingConnection(serverId: string): MCPConnection | undefined { + return this.connections.get(serverId); + } - if (!serverConfig) { - throw new Error(`[MCPStore] No config found for ${serverName}, cannot reconnect`); + /** + * Get server instructions from health check results (for display before active connection). + * Useful for showing instructions in settings UI. + */ + getHealthCheckInstructions(): Array<{ + serverId: string; + serverTitle?: string; + instructions: string; + }> { + const results: Array<{ serverId: string; serverTitle?: string; instructions: string }> = []; + + for (const [serverId, state] of Object.entries(this.health.checks)) { + if (state.status === HealthCheckStatus.SUCCESS && state.instructions) { + results.push({ + instructions: state.instructions, + serverId, + serverTitle: state.serverInfo?.title || state.serverInfo?.name + }); + } } - // Disconnect stale connection (clears old transport + session ID) - const oldConnection = this.connections.get(serverName); - - if (oldConnection) { - await MCPService.disconnect(oldConnection).catch(console.warn); - this.connections.delete(serverName); - } - - console.log(`[MCPStore][${serverName}] Session expired, reconnecting with fresh session...`); - - const listChangedHandlers = this.createListChangedHandlers(serverName); - const connection = await MCPService.connect( - serverName, - serverConfig, - DEFAULT_MCP_CONFIG.clientInfo, - DEFAULT_MCP_CONFIG.capabilities, - (phase) => { - if (phase === MCPConnectionPhase.DISCONNECTED) { - console.log(`[MCPStore][${serverName}] Connection lost, starting auto-reconnect`); - this.autoReconnect(serverName); - } - }, - listChangedHandlers - ); - - // Replace connection and rebuild tool index for this server - this.connections.set(serverName, connection); - for (const tool of connection.tools) { - this.toolsIndex.set(tool.name, serverName); - } - - console.log(`[MCPStore][${serverName}] Session recovered successfully`); - } + return results; + } /** - * Auto-reconnect to a server with exponential backoff. - * Continues indefinitely until successful. - * - * Race-condition safety: when the phase callback fires a DISCONNECTED event - * while we are still inside this function (e.g., the server drops right after - * a successful connect()), a naive inner `autoReconnect()` call would be - * swallowed by the `reconnectingServers` guard, leaving the server - * permanently disconnected once the outer call exits. We solve this by - * deferring the new reconnection via the `needsReconnect` flag: the flag is - * set inside the phase callback and honoured in the `finally` block after - * the guard entry has been removed. + * Health checks live in MCPHealthCheckManager; these delegate so + * consumers keep a single entry point. */ - private async autoReconnect(serverName: string): Promise { - // Guard against concurrent reconnections - if (this.reconnectingServers.has(serverName)) { - console.log(`[MCPStore][${serverName}] Reconnection already in progress, skipping`); + getHealthCheckState(serverId: string): HealthCheckState { + return this.health.getState(serverId); + } - return; - } + async getPrompt( + serverName: string, + promptName: string, + args?: Record + ): Promise { + const connection = this.connections.get(serverName); - const serverConfig = this.serverConfigs.get(serverName); + if (!connection) throw new Error(`Server "${serverName}" not found for prompt "${promptName}"`); - if (!serverConfig) { - console.error(`[MCPStore] No config found for ${serverName}, cannot reconnect`); + return MCPService.getPrompt(connection, promptName, args); + } - return; - } + async getPromptCompletions( + serverName: string, + promptName: string, + argumentName: string, + argumentValue: string + ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { + const connection = this.connections.get(serverName); - this.reconnectingServers.add(serverName); - let backoff = MCP_RECONNECT.INITIAL_DELAY; - // Flag set by the phase callback when a DISCONNECTED event fires while - // reconnectingServers still holds this server (see JSDoc above). - let needsReconnect = false; + if (!connection) { + console.warn(`[MCPStore] Server "${serverName}" is not connected`); - try { - while (true) { - await new Promise((resolve) => setTimeout(resolve, backoff)); + return null; + } - console.log(`[MCPStore][${serverName}] Auto-reconnecting...`); + if (!connection.serverCapabilities?.completions) { + return null; + } - try { - // Per-attempt timeout: reject if the server doesn't respond in time, - // then fall through to backoff logic as with any other failure. - const timeoutPromise = new Promise((_, reject) => - setTimeout( - () => - reject( - new Error( - `Reconnect attempt timed out after ${MCP_RECONNECT.ATTEMPT_TIMEOUT_MS}ms` - ) - ), - MCP_RECONNECT.ATTEMPT_TIMEOUT_MS - ) - ); + return MCPService.complete( + connection, + { name: promptName, type: MCPRefType.PROMPT }, + { name: argumentName, value: argumentValue } + ); + } - needsReconnect = false; - const listChangedHandlers = this.createListChangedHandlers(serverName); - const connectPromise = MCPService.connect( - serverName, - serverConfig, - DEFAULT_MCP_CONFIG.clientInfo, - DEFAULT_MCP_CONFIG.capabilities, - (phase) => { - if (phase === MCPConnectionPhase.DISCONNECTED) { - if (this.reconnectingServers.has(serverName)) { - // Reconnect loop is active; defer to after it exits. - needsReconnect = true; - } else { - console.log( - `[MCPStore][${serverName}] Connection lost, restarting auto-reconnect` - ); - this.autoReconnect(serverName); - } - } - }, - listChangedHandlers - ); - const connection = await Promise.race([connectPromise, timeoutPromise]); + /** + * Request timeout in milliseconds, read live from the global setting + * so a change in Settings applies to every server immediately. + */ + getRequestTimeoutMs(): number { + const seconds = + Number(settingsStore.config.mcpRequestTimeoutSeconds) || + DEFAULT_MCP_CONFIG.requestTimeoutSeconds; - // Replace old connection with new one - this.connections.set(serverName, connection); + return Math.round(seconds * 1000); + } - // Rebuild tool index for this server - for (const tool of connection.tools) { - this.toolsIndex.set(tool.name, serverName); - } + /** + * Get completions for a resource template argument. + * Uses the MCP Completion API with ref/resource. + */ + async getResourceCompletions( + serverName: string, + uriTemplate: string, + argumentName: string, + argumentValue: string + ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { + const connection = this.connections.get(serverName); - console.log(`[MCPStore][${serverName}] Reconnected successfully`); + if (!connection) { + console.warn(`[MCPStore] Server "${serverName}" is not connected`); - break; - } catch (error) { - console.warn(`[MCPStore][${serverName}] Reconnection failed:`, error); - backoff = Math.min(backoff * MCP_RECONNECT.BACKOFF_MULTIPLIER, MCP_RECONNECT.MAX_DELAY); - } - } - } finally { - this.reconnectingServers.delete(serverName); + return null; + } - // If the phase callback signalled a disconnect while this function held - // the guard, kick off a fresh reconnect now that the guard is released. - if (needsReconnect) { - console.log( - `[MCPStore][${serverName}] Deferred disconnect detected, restarting auto-reconnect` - ); - this.autoReconnect(serverName); - } + if (!connection.serverCapabilities?.completions) { + return null; } + + return MCPService.complete( + connection, + { type: MCPRefType.RESOURCE, uri: uriTemplate }, + { name: argumentName, value: argumentValue } + ); } - getToolNames(): string[] { - return Array.from(this.toolsIndex.keys()); + /** + * Get formatted resource context for chat. + */ + getResourceContextForChat(): string { + return mcpResourceStore.formatAttachmentsForContext(); } - hasTool(toolName: string): boolean { - return this.toolsIndex.has(toolName); + getServerById(serverId: string): MCPServerSettingsEntry | undefined { + return this.getServers().find((s) => s.id === serverId); } - getToolServer(toolName: string): string | undefined { - return this.toolsIndex.get(toolName); + /** + * Get display name for an MCP server by its ID. + * Falls back to the server ID if server is not found. + */ + getServerDisplayName(serverId: string): string { + const server = this.getServerById(serverId); + + return server ? this.getServerLabel(server) : serverId; } /** - * Resolve which configured MCP server owns a given tool name. Looks at - * active connections first (fast path), then falls back to per-server - * health-check data so server-side MCP proxies (where llama-server - * executes MCP tools but the browser does not hold a direct connection) - * still resolve tool names to their owning server. + * Get icon URL for an MCP server by its ID. + * Returns the best icon from the MCP server's `icons` array + * (see MCP spec: spec.modelcontextprotocol.io). + * Returns null if no icon is available. */ - findServerForTool(toolName: string): string | undefined { - const fromIndex = this.toolsIndex.get(toolName); + getServerFavicon(serverId: string): string | null { + const server = this.getServerById(serverId); - if (fromIndex) return fromIndex; + if (!server) { + return null; + } - for (const server of this.getServers()) { - const health = this._healthChecks[server.id]; + const isDark = mode.current === ColorMode.DARK; + const healthState = this.health.getState(serverId); - if (!health || health.status !== HealthCheckStatus.SUCCESS) continue; + if (healthState.status === HealthCheckStatus.SUCCESS && healthState.serverInfo?.icons) { + const mcpIconUrl = getMcpIconUrl(healthState.serverInfo.icons, isDark); - if (health.tools.some((tool) => tool.name === toolName)) { - return server.id; + if (mcpIconUrl) { + return mcpIconUrl; } } - return undefined; + return getMcpServerFaviconFallback(server.url); } /** @@ -1051,26 +702,126 @@ class MCPStore { return this.getServerFavicon(serverId); } - hasPromptsSupport(): boolean { - for (const connection of this.connections.values()) { - if (connection.serverCapabilities?.prompts) { - return true; + /** + * Get aggregated server instructions from all connected servers. + * Returns an array of { serverName, serverTitle, instructions } objects. + */ + getServerInstructions(): Array<{ + serverName: string; + serverTitle?: string; + instructions: string; + }> { + const results: Array<{ serverName: string; serverTitle?: string; instructions: string }> = []; + + for (const [serverName, connection] of this.connections) { + if (connection.instructions) { + results.push({ + instructions: connection.instructions, + serverName, + serverTitle: connection.serverInfo?.title || connection.serverInfo?.name + }); } } - return false; + return results; } - /** - * Check if any enabled server with successful health check supports prompts. - * Uses health check state since servers may not have active connections until - * the user actually sends a message or uses prompts. - * @param perChatOverrides - Per-chat server overrides to filter by enabled servers. - * If provided (even empty array), only checks enabled servers. - * If undefined, falls back to each server's own `enabled` flag. - */ - hasPromptsCapability(perChatOverrides?: McpServerOverride[]): boolean { - let enabledServerIds: Set; + getServerLabel(server: MCPServerDisplayInfo): string { + return getMcpServerLabel(server, this.getServers(), this.health.checks); + } + + getServers(): MCPServerSettingsEntry[] { + const raw = settingsStore.config.mcpServers; + + // cache the parse: the config string rarely changes and getServers is + // called from hot paths (per-tool display lookups, capability checks) + if (this.serversCache && this.serversCache.raw === raw) { + return this.serversCache.servers; + } + + const servers = parseMcpServerSettings(raw); + + this.serversCache = { raw, servers }; + + return servers; + } + + getServersStatus(): ServerStatus[] { + const statuses: ServerStatus[] = []; + + for (const [name, connection] of this.connections) { + statuses.push({ + error: undefined, + isConnected: true, + name, + toolCount: connection.tools.length + }); + } + + return statuses; + } + + /** + * Get list of enabled servers that support resources. + * Checks active connections first, then health check state as fallback. + */ + getServersWithResources(): string[] { + const enabledServerIds = new Set( + this.getServers() + .filter((s) => s.enabled) + .map((s) => s.id) + ); + const servers: string[] = []; + + for (const [name, connection] of this.connections) { + if (!enabledServerIds.has(name)) continue; + + if (MCPService.supportsResources(connection) && !servers.includes(name)) { + servers.push(name); + } + } + + // Also check health check states for servers not yet connected + for (const [serverId, state] of Object.entries(this.health.checks)) { + if (!enabledServerIds.has(serverId)) continue; + + if ( + !servers.includes(serverId) && + state.status === HealthCheckStatus.SUCCESS && + state.capabilities?.server?.resources !== undefined + ) { + servers.push(serverId); + } + } + + return servers; + } + + getToolNames(): string[] { + return Array.from(this.toolsIndex.keys()); + } + + getToolServer(toolName: string): string | undefined { + return this.toolsIndex.get(toolName); + } + + hasAvailableServers(): boolean { + return parseMcpServerSettings(settingsStore.config.mcpServers).some( + (s) => s.enabled && s.url.trim() + ); + } + + hasEnabledServers(perChatOverrides?: McpServerOverride[]): boolean { + return Boolean(this.buildMcpClientConfig(settingsStore.config, perChatOverrides)); + } + + /** + * Check if any enabled server with successful health check supports prompts. + * Uses health check state since servers may not have active connections until + * the user actually sends a message or uses prompts. + */ + hasPromptsCapability(perChatOverrides?: McpServerOverride[]): boolean { + let enabledServerIds: Set; if (perChatOverrides !== undefined) { enabledServerIds = new Set(perChatOverrides.filter((o) => o.enabled).map((o) => o.serverId)); @@ -1086,7 +837,7 @@ class MCPStore { return false; } - for (const [serverId, state] of Object.entries(this._healthChecks)) { + for (const [serverId, state] of Object.entries(this.health.checks)) { if (!enabledServerIds.has(serverId)) continue; if ( @@ -1108,185 +859,134 @@ class MCPStore { return false; } - async getAllPrompts(): Promise { - const results: MCPPromptInfo[] = []; - - for (const [serverName, connection] of this.connections) { - if (!connection.serverCapabilities?.prompts) continue; - - const prompts = await MCPService.listPrompts(connection); - - for (const prompt of prompts) { - results.push({ - arguments: prompt.arguments?.map((arg) => ({ - description: arg.description, - name: arg.name, - required: arg.required - })), - description: prompt.description, - name: prompt.name, - serverName, - title: prompt.title - }); + hasPromptsSupport(): boolean { + for (const connection of this.connections.values()) { + if (connection.serverCapabilities?.prompts) { + return true; } } - return results; - } - - async getPrompt( - serverName: string, - promptName: string, - args?: Record - ): Promise { - const connection = this.connections.get(serverName); - - if (!connection) throw new Error(`Server "${serverName}" not found for prompt "${promptName}"`); - - return MCPService.getPrompt(connection, promptName, args); + return false; } - async executeTool(toolCall: MCPToolCall, signal?: AbortSignal): Promise { - const toolName = toolCall.function.name; - const serverName = this.toolsIndex.get(toolName); + /** + * Check if any enabled server with successful health check supports resources. + * Uses health check state since servers may not have active connections until + * the user actually sends a message or uses prompts. + */ + hasResourcesCapability(perChatOverrides?: McpServerOverride[]): boolean { + let enabledServerIds: Set; - if (!serverName) throw new Error(`Unknown tool: ${toolName}`); + if (perChatOverrides !== undefined) { + enabledServerIds = new Set(perChatOverrides.filter((o) => o.enabled).map((o) => o.serverId)); + } else { + enabledServerIds = new Set( + this.getServers() + .filter((s) => s.enabled) + .map((s) => s.id) + ); + } - const connection = this.connections.get(serverName); + if (enabledServerIds.size === 0) { + return false; + } - if (!connection) throw new Error(`Server "${serverName}" is not connected`); + for (const [serverId, state] of Object.entries(this.health.checks)) { + if (!enabledServerIds.has(serverId)) continue; - const args = this.parseToolArguments(toolCall.function.arguments); + if ( + state.status === HealthCheckStatus.SUCCESS && + state.capabilities?.server?.resources !== undefined + ) { + return true; + } + } - try { - return await MCPService.callTool(connection, { arguments: args, name: toolName }, signal); - } catch (error) { - // Session expired (server restarted) - reconnect and retry once - if (MCPService.isSessionExpiredError(error)) { - await this.reconnectServer(serverName); + for (const [serverName, connection] of this.connections) { + if (!enabledServerIds.has(serverName)) continue; - const newConnection = this.connections.get(serverName); + if (MCPService.supportsResources(connection)) { + return true; + } + } - if (!newConnection) throw new Error(`Failed to reconnect to "${serverName}"`); + return false; + } - return MCPService.callTool(newConnection, { arguments: args, name: toolName }, signal); + /** + * Check if any connected server has instructions. + */ + hasServerInstructions(): boolean { + for (const connection of this.connections.values()) { + if (connection.instructions) { + return true; } - - throw error; } + + return false; } - async executeToolByName( - toolName: string, - args: Record, - signal?: AbortSignal - ): Promise { - const serverName = this.toolsIndex.get(toolName); + hasTool(toolName: string): boolean { + return this.toolsIndex.has(toolName); + } - if (!serverName) throw new Error(`Unknown tool: ${toolName}`); + /** + * Promote a health check connection to an active connection. + * This avoids the need to reconnect when the server is needed for agentic flows. + */ + promoteHealthCheckToConnection(serverId: string, connection: MCPConnection): void { + this.indexServerTools(serverId, connection.tools); - const connection = this.connections.get(serverName); + this.connections.set(serverId, connection); - if (!connection) throw new Error(`Server "${serverName}" is not connected`); + this.updateState({ + connectedServers: Array.from(this.connections.keys()), + toolCount: this.toolsIndex.size + }); + } - try { - return await MCPService.callTool(connection, { arguments: args, name: toolName }, signal); - } catch (error) { - if (MCPService.isSessionExpiredError(error)) { - await this.reconnectServer(serverName); + /** + * Read resource content from a server. + * Caches the result in mcpResourceStore. + */ + async readResource(uri: string): Promise { + const cached = mcpResourceStore.getCachedContent(uri); - const newConnection = this.connections.get(serverName); + if (cached) { + return cached.content; + } - if (!newConnection) throw new Error(`Failed to reconnect to "${serverName}"`); + // Find which server has this resource + const serverName = mcpResourceStore.findServerForUri(uri); - return MCPService.callTool(newConnection, { arguments: args, name: toolName }, signal); - } + if (!serverName) { + console.error(`[MCPStore] No server found for resource URI: ${uri}`); - throw error; + return null; } - } - private parseToolArguments(args: string | Record): Record { - if (typeof args === 'string') { - const trimmed = args.trim(); + const connection = this.connections.get(serverName); - if (trimmed === '') { - return {}; - } + if (!connection) { + console.error(`[MCPStore] No connection found for server: ${serverName}`); - try { - const parsed = JSON.parse(trimmed); + return null; + } - if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) - throw new Error( - `Tool arguments must be an object, got ${Array.isArray(parsed) ? 'array' : typeof parsed}` - ); + try { + const result = await MCPService.readResource(connection, uri); + const resourceInfo = mcpResourceStore.findResourceByUri(uri); - return parsed as Record; - } catch (error) { - throw new Error(`Failed to parse tool arguments as JSON: ${(error as Error).message}`); + if (resourceInfo) { + mcpResourceStore.cacheResourceContent(resourceInfo, result.contents); } - } - if (typeof args === 'object' && args !== null && !Array.isArray(args)) { - return args; - } - - throw new Error(`Invalid tool arguments type: ${typeof args}`); - } - - async getPromptCompletions( - serverName: string, - promptName: string, - argumentName: string, - argumentValue: string - ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { - const connection = this.connections.get(serverName); - - if (!connection) { - console.warn(`[MCPStore] Server "${serverName}" is not connected`); - - return null; - } - - if (!connection.serverCapabilities?.completions) { - return null; - } - - return MCPService.complete( - connection, - { name: promptName, type: MCPRefType.PROMPT }, - { name: argumentName, value: argumentValue } - ); - } - - /** - * Get completions for a resource template argument. - * Uses the MCP Completion API with ref/resource. - */ - async getResourceCompletions( - serverName: string, - uriTemplate: string, - argumentName: string, - argumentValue: string - ): Promise<{ values: string[]; total?: number; hasMore?: boolean } | null> { - const connection = this.connections.get(serverName); - - if (!connection) { - console.warn(`[MCPStore] Server "${serverName}" is not connected`); - - return null; - } + return result.contents; + } catch (error) { + console.error(`[MCPStore] Failed to read resource ${uri}:`, error); - if (!connection.serverCapabilities?.completions) { return null; } - - return MCPService.complete( - connection, - { type: MCPRefType.RESOURCE, uri: uriTemplate }, - { name: argumentName, value: argumentValue } - ); } /** @@ -1313,21 +1013,51 @@ class MCPStore { } } - private parseHeaders(headersJson?: string): Record | undefined { - if (!headersJson?.trim()) { - return undefined; - } + /** Store a server config so auto-reconnect can rebuild the session. */ + registerServerConfig(name: string, config: MCPServerConfig): void { + this.serverConfigs.set(name, config); + } - try { - const parsed = JSON.parse(headersJson); + /** + * Release a connection reference. + * By default, keeps connections alive for reuse (shutdownIfUnused=false). + * MCP spec encourages long-lived sessions to avoid reconnection overhead. + */ + async releaseConnection(shutdownIfUnused = false): Promise { + this.activeFlowCount = Math.max(0, this.activeFlowCount - 1); - if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) - return parsed as Record; - } catch { - console.warn('[MCPStore] Failed to parse custom headers JSON:', headersJson); + if (shutdownIfUnused && this.activeFlowCount === 0) { + await this.shutdown(); } + } - return undefined; + /** + * Drop a connection without disconnecting, e.g. when a health check finds + * it stale and recreates it. + */ + removeConnection(serverId: string): void { + this.connections.delete(serverId); + } + + /** + * Remove a resource attachment from chat context. + */ + removeResourceAttachment(attachmentId: string): void { + mcpResourceStore.removeAttachment(attachmentId); + } + + removeServer(id: string): void { + const servers = this.getServers(); + + settingsStore.updateConfig( + SETTINGS_KEYS.MCP_SERVERS, + JSON.stringify(servers.filter((s) => s.id !== id)) + ); + this.clearHealthCheck(id); + } + + async runHealthCheck(server: HealthCheckParams, promoteToActive = false): Promise { + return this.health.run(server, promoteToActive); } async runHealthChecksForServers( @@ -1340,644 +1070,467 @@ class MCPStore { skipIfChecked = true, promoteToActive = false ): Promise { - const serversToCheck = skipIfChecked - ? servers.filter((s) => !this.hasHealthCheck(s.id) && s.url.trim()) - : servers.filter((s) => s.url.trim()); + return this.health.runForServers(servers, skipIfChecked, promoteToActive); + } - if (serversToCheck.length === 0) { - return; + async shutdown(): Promise { + if (this.initPromise) { + await this.initPromise.catch(() => {}); + this.initPromise = null; } - const BATCH_SIZE = 5; - - for (let i = 0; i < serversToCheck.length; i += BATCH_SIZE) { - const batch = serversToCheck.slice(i, i + BATCH_SIZE); - - await Promise.allSettled(batch.map((server) => this.runHealthCheck(server, promoteToActive))); + if (this.connections.size === 0) { + return; } - } - /** - * Check if a server already has an active connection that can be reused. - * Returns the existing connection if available. - */ - getExistingConnection(serverId: string): MCPConnection | undefined { - return this.connections.get(serverId); + await Promise.all( + Array.from(this.connections.values()).map((conn) => + MCPService.disconnect(conn).catch((error) => + console.warn(`[MCPStore] Error disconnecting ${conn.serverName}:`, error) + ) + ) + ); + + this.connections.clear(); + this.toolsIndex.clear(); + this.serverConfigs.clear(); + this.configSignature = null; + this.updateState({ + connectedServers: [], + error: null, + isInitializing: false, + toolCount: 0 + }); } /** - * Run a health check for a server. - * If the server already has an active connection, reuses it instead of creating a new one. - * If promoteToActive is true and server is enabled, the connection will be kept - * and promoted to an active connection instead of being disconnected. + * Subscribe to resource updates. */ - async runHealthCheck(server: HealthCheckParams, promoteToActive = false): Promise { - // Check if we already have an active connection for this server - const existingConnection = this.connections.get(server.id); - - if (existingConnection) { - // Reuse existing connection - just refresh tools list - try { - const tools = await MCPService.listTools(existingConnection); - const capabilities = this.#buildCapabilitiesInfo( - existingConnection.serverCapabilities, - existingConnection.clientCapabilities - ); + async subscribeToResource(uri: string): Promise { + const serverName = mcpResourceStore.findServerForUri(uri); - this.updateHealthCheck(server.id, { - capabilities, - connectionTimeMs: existingConnection.connectionTimeMs, - instructions: existingConnection.instructions, - logs: [], - protocolVersion: existingConnection.protocolVersion, - serverInfo: existingConnection.serverInfo, - status: HealthCheckStatus.SUCCESS, - tools: tools.map((tool) => ({ - description: tool.description, - name: tool.name, - title: tool.title - })), - transportType: existingConnection.transportType - }); + if (!serverName) { + console.error(`[MCPStore] No server found for resource URI: ${uri}`); - return; - } catch (error) { - console.warn( - `[MCPStore] Failed to reuse connection for ${server.id}, creating new one:`, - error - ); - // Connection may be stale, remove it and create new one - this.connections.delete(server.id); - } + return false; } - const trimmedUrl = server.url.trim(); - const logs: MCPConnectionLog[] = []; - - let currentPhase: MCPConnectionPhase = MCPConnectionPhase.IDLE; + const connection = this.connections.get(serverName); - if (!trimmedUrl) { - this.updateHealthCheck(server.id, { - logs: [], - message: 'Please enter a server URL first.', - status: HealthCheckStatus.ERROR - }); + if (!connection) { + console.error(`[MCPStore] No connection found for server: ${serverName}`); - return; + return false; } - this.updateHealthCheck(server.id, { - logs: [], - phase: MCPConnectionPhase.TRANSPORT_CREATING, - status: HealthCheckStatus.CONNECTING - }); - - const timeoutMs = this.#requestTimeoutMs(); - const headers = this.parseHeaders(server.headers); + if (!MCPService.supportsResourceSubscriptions(connection)) { + return false; + } try { - const serverConfig: MCPServerConfig = { - handshakeTimeoutMs: DEFAULT_MCP_CONFIG.connectionTimeoutMs, - headers, - requestTimeoutMs: timeoutMs, - transport: detectMcpTransportFromUrl(trimmedUrl), - url: trimmedUrl, - useProxy: server.useProxy - }; - - // Store config for reconnection - this.serverConfigs.set(server.id, serverConfig); - - const connection = await MCPService.connect( - server.id, - serverConfig, - DEFAULT_MCP_CONFIG.clientInfo, - DEFAULT_MCP_CONFIG.capabilities, - (phase, log) => { - currentPhase = phase; - logs.push(log); - this.updateHealthCheck(server.id, { - logs: [...logs], - phase, - status: HealthCheckStatus.CONNECTING - }); - - // Handle WebSocket disconnection - if (phase === MCPConnectionPhase.DISCONNECTED && promoteToActive) { - console.log( - `[MCPStore][${server.id}] Connection lost during health check, starting auto-reconnect` - ); - this.autoReconnect(server.id); - } - } - ); - const tools = connection.tools.map((tool) => ({ - description: tool.description, - name: tool.name, - title: tool.title - })); - const capabilities = this.#buildCapabilitiesInfo( - connection.serverCapabilities, - connection.clientCapabilities - ); - - this.updateHealthCheck(server.id, { - capabilities, - connectionTimeMs: connection.connectionTimeMs, - instructions: connection.instructions, - logs, - protocolVersion: connection.protocolVersion, - serverInfo: connection.serverInfo, - status: HealthCheckStatus.SUCCESS, - tools, - transportType: connection.transportType - }); + await MCPService.subscribeResource(connection, uri); + mcpResourceStore.addSubscription(uri, serverName); - // Promote to active connection or disconnect - if (promoteToActive && server.enabled) { - this.promoteHealthCheckToConnection(server.id, connection); - } else { - await MCPService.disconnect(connection); - } + return true; } catch (error) { - const message = error instanceof Error ? error.message : 'Unknown error occurred'; - - if (logs.at(-1)?.phase !== MCPConnectionPhase.ERROR) { - logs.push({ - level: MCPLogLevel.ERROR, - message: `Connection failed: ${message}`, - phase: MCPConnectionPhase.ERROR, - timestamp: new Date() - }); - } + console.error(`[MCPStore] Failed to subscribe to resource ${uri}:`, error); - this.updateHealthCheck(server.id, { - logs, - message, - phase: currentPhase, - status: HealthCheckStatus.ERROR - }); + return false; } } /** - * Promote a health check connection to an active connection. - * This avoids the need to reconnect when the server is needed for agentic flows. + * Unsubscribe from resource updates. */ - private promoteHealthCheckToConnection(serverId: string, connection: MCPConnection): void { - // Register tools from the connection - for (const tool of connection.tools) { - if (this.toolsIndex.has(tool.name)) { - console.warn( - `[MCPStore] Tool name conflict during promotion: "${tool.name}" exists in "${this.toolsIndex.get(tool.name)}" and "${serverId}". Using tool from "${serverId}".` - ); - } - - this.toolsIndex.set(tool.name, serverId); - } - - // Add to active connections - this.connections.set(serverId, connection); - - // Update state - this.updateState({ - connectedServers: Array.from(this.connections.keys()), - toolCount: this.toolsIndex.size - }); - } + async unsubscribeFromResource(uri: string): Promise { + const serverName = mcpResourceStore.findServerForUri(uri); - getServersStatus(): ServerStatus[] { - const statuses: ServerStatus[] = []; + if (!serverName) { + console.error(`[MCPStore] No server found for resource URI: ${uri}`); - for (const [name, connection] of this.connections) { - statuses.push({ - error: undefined, - isConnected: true, - name, - toolCount: connection.tools.length - }); + return false; } - return statuses; - } + const connection = this.connections.get(serverName); - /** - * Get aggregated server instructions from all connected servers. - * Returns an array of { serverName, serverTitle, instructions } objects. - */ - getServerInstructions(): Array<{ - serverName: string; - serverTitle?: string; - instructions: string; - }> { - const results: Array<{ serverName: string; serverTitle?: string; instructions: string }> = []; + if (!connection) { + console.error(`[MCPStore] No connection found for server: ${serverName}`); - for (const [serverName, connection] of this.connections) { - if (connection.instructions) { - results.push({ - instructions: connection.instructions, - serverName, - serverTitle: connection.serverInfo?.title || connection.serverInfo?.name - }); - } + return false; } - return results; - } + try { + await MCPService.unsubscribeResource(connection, uri); + mcpResourceStore.removeSubscription(uri); - /** - * Get server instructions from health check results (for display before active connection). - * Useful for showing instructions in settings UI. - */ - getHealthCheckInstructions(): Array<{ - serverId: string; - serverTitle?: string; - instructions: string; - }> { - const results: Array<{ serverId: string; serverTitle?: string; instructions: string }> = []; + return true; + } catch (error) { + console.error(`[MCPStore] Failed to unsubscribe from resource ${uri}:`, error); - for (const [serverId, state] of Object.entries(this._healthChecks)) { - if (state.status === HealthCheckStatus.SUCCESS && state.instructions) { - results.push({ - instructions: state.instructions, - serverId, - serverTitle: state.serverInfo?.title || state.serverInfo?.name - }); - } + return false; } - - return results; } - /** - * Check if any connected server has instructions. - */ - hasServerInstructions(): boolean { - for (const connection of this.connections.values()) { - if (connection.instructions) { - return true; - } - } + updateServer(id: string, updates: Partial): void { + const servers = this.getServers(); - return false; + settingsStore.updateConfig( + SETTINGS_KEYS.MCP_SERVERS, + JSON.stringify( + servers.map((server) => (server.id === id ? { ...server, ...updates } : server)) + ) + ); } /** - * - * - * Resources Operations - * - * - */ - - /** - * Check if any enabled server with successful health check supports resources. - * Uses health check state since servers may not have active connections until - * the user actually sends a message or uses prompts. - * @param perChatOverrides - Per-chat server overrides to filter by enabled servers. - * If provided (even empty array), only checks enabled servers. - * If undefined, falls back to each server's own `enabled` flag. + * Builds MCP client configuration from settings. */ - hasResourcesCapability(perChatOverrides?: McpServerOverride[]): boolean { - let enabledServerIds: Set; - - if (perChatOverrides !== undefined) { - enabledServerIds = new Set(perChatOverrides.filter((o) => o.enabled).map((o) => o.serverId)); - } else { - enabledServerIds = new Set( - this.getServers() - .filter((s) => s.enabled) - .map((s) => s.id) - ); - } - - if (enabledServerIds.size === 0) { - return false; - } - - for (const [serverId, state] of Object.entries(this._healthChecks)) { - if (!enabledServerIds.has(serverId)) continue; - - if ( - state.status === HealthCheckStatus.SUCCESS && - state.capabilities?.server?.resources !== undefined - ) { - return true; - } - } - - for (const [serverName, connection] of this.connections) { - if (!enabledServerIds.has(serverName)) continue; + private buildMcpClientConfig( + cfg: SettingsConfigType, + perChatOverrides?: McpServerOverride[] + ): MCPClientConfig | undefined { + const rawServers = parseMcpServerSettings(cfg.mcpServers); - if (MCPService.supportsResources(connection)) { - return true; - } + if (!rawServers.length) { + return undefined; } - return false; - } + const servers: Record = {}; - /** - * Get list of enabled servers that support resources. - * Checks active connections first, then health check state as fallback. - */ - getServersWithResources(): string[] { - const enabledServerIds = new Set( - this.getServers() - .filter((s) => s.enabled) - .map((s) => s.id) - ); - const servers: string[] = []; + for (const [index, entry] of rawServers.entries()) { + if (!this.checkServerEnabled(entry, perChatOverrides)) continue; - // Check active connections - for (const [name, connection] of this.connections) { - if (!enabledServerIds.has(name)) continue; + const normalized = this.buildServerConfig(entry); - if (MCPService.supportsResources(connection) && !servers.includes(name)) { - servers.push(name); - } + if (normalized) servers[this.generateServerId(entry.id, index)] = normalized; } - // Also check health check states for servers not yet connected - for (const [serverId, state] of Object.entries(this._healthChecks)) { - if (!enabledServerIds.has(serverId)) continue; - - if ( - !servers.includes(serverId) && - state.status === HealthCheckStatus.SUCCESS && - state.capabilities?.server?.resources !== undefined - ) { - servers.push(serverId); - } + if (Object.keys(servers).length === 0) { + return undefined; } - return servers; + return { + capabilities: DEFAULT_MCP_CONFIG.capabilities, + clientInfo: DEFAULT_MCP_CONFIG.clientInfo, + protocolVersion: DEFAULT_MCP_CONFIG.protocolVersion, + requestTimeoutMs: this.getRequestTimeoutMs(), + servers + }; } /** - * Fetch resources from all connected servers that support them. - * Updates mcpResourceStore with the results. - * @param forceRefresh - If true, bypass cache and fetch fresh data + * Builds server configuration from a settings entry. */ - async fetchAllResources(forceRefresh: boolean = false): Promise { - const serversWithResources = this.getServersWithResources(); - - if (serversWithResources.length === 0) { - return; + private buildServerConfig( + entry: MCPServerSettingsEntry, + connectionTimeoutMs = DEFAULT_MCP_CONFIG.connectionTimeoutMs + ): MCPServerConfig | undefined { + if (!entry?.url) { + return undefined; } - // Check if we have cached resources and they're recent (unless force refresh) - if (!forceRefresh) { - const allServersCached = serversWithResources.every((serverName) => { - const serverRes = mcpResourceStore.getServerResources(serverName); - - if (!serverRes || !serverRes.lastFetched) { - return false; - } - - // Cache is valid for 5 minutes - const age = Date.now() - serverRes.lastFetched.getTime(); - - return age < CACHE.DEFAULT_TTL_MS; - }); + let headers: Record | undefined; - if (allServersCached) { - console.log('[MCPStore] Using cached resources'); + if (entry.headers) { + try { + const parsed = JSON.parse(entry.headers); - return; + if (typeof parsed === 'object' && parsed !== null && !Array.isArray(parsed)) + headers = parsed as Record; + } catch { + console.warn('[MCP] Failed to parse custom headers JSON:', entry.headers); } } - mcpResourceStore.setLoading(true); - - try { - await Promise.all( - serversWithResources.map((serverName) => this.fetchServerResources(serverName)) - ); - } finally { - mcpResourceStore.setLoading(false); - } + return { + handshakeTimeoutMs: connectionTimeoutMs, + headers, + requestTimeoutMs: this.getRequestTimeoutMs(), + transport: detectMcpTransportFromUrl(entry.url), + url: entry.url, + useProxy: entry.useProxy + }; } /** - * Fetch resources from a specific server. - * Updates mcpResourceStore with the results. + * Checks if a server is enabled for a given chat. + * A per-chat override wins when present; a server without one resolves + * to its own `enabled` flag in `mcpServers`. */ - async fetchServerResources(serverName: string): Promise { - const connection = this.connections.get(serverName); - - if (!connection) { - console.warn(`[MCPStore] No connection found for server: ${serverName}`); - - return; - } + private checkServerEnabled( + server: MCPServerSettingsEntry, + perChatOverrides?: McpServerOverride[] + ): boolean { + const override = perChatOverrides?.find((o) => o.serverId === server.id); - if (!MCPService.supportsResources(connection)) { - return; - } + return override?.enabled ?? server.enabled; + } - mcpResourceStore.setServerLoading(serverName, true); + private createListChangedHandlers(serverName: string): ListChangedHandlers { + return { + prompts: { + onChanged: (error: Error | null) => { + if (error) { + console.warn(`[MCPStore][${serverName}] Prompts list changed error:`, error); - try { - const [resources, templates] = await Promise.all([ - MCPService.listAllResources(connection), - MCPService.listAllResourceTemplates(connection) - ]); + return; + } + } + }, + tools: { + onChanged: (error: Error | null, tools: Tool[] | null) => { + if (error) { + console.warn(`[MCPStore][${serverName}] Tools list changed error:`, error); - mcpResourceStore.setServerResources(serverName, resources, templates); - } catch (error) { - const message = error instanceof Error ? error.message : String(error); + return; + } - mcpResourceStore.setServerError(serverName, message); - console.error(`[MCPStore][${serverName}] Failed to fetch resources:`, error); - } + this.handleToolsListChanged(serverName, tools ?? []); + } + } + }; } - /** - * Read resource content from a server. - * Caches the result in mcpResourceStore. - */ - async readResource(uri: string): Promise { - // Check cache first - const cached = mcpResourceStore.getCachedContent(uri); + private async doInitialize( + signature: string, + mcpConfig: MCPClientConfig, + serverEntries: [string, MCPClientConfig['servers'][string]][] + ): Promise { + const clientInfo = mcpConfig.clientInfo ?? DEFAULT_MCP_CONFIG.clientInfo; + const capabilities = mcpConfig.capabilities ?? DEFAULT_MCP_CONFIG.capabilities; + const results = await Promise.allSettled( + serverEntries.map(async ([name, serverConfig]) => { + this.serverConfigs.set(name, serverConfig); - if (cached) { - return cached.content; - } + const listChangedHandlers = this.createListChangedHandlers(name); + const connection = await MCPService.connect( + name, + serverConfig, + clientInfo, + capabilities, + (phase) => { + if (phase === MCPConnectionPhase.DISCONNECTED) { + console.log(`[MCPStore][${name}] Connection lost, starting auto-reconnect`); + this.autoReconnect(name); + } + }, + listChangedHandlers + ); - // Find which server has this resource - const serverName = mcpResourceStore.findServerForUri(uri); + return { connection, name }; + }) + ); - if (!serverName) { - console.error(`[MCPStore] No server found for resource URI: ${uri}`); + if (this.configSignature !== signature) { + for (const result of results) { + if (result.status === 'fulfilled') + await MCPService.disconnect(result.value.connection).catch(console.warn); + } - return null; + return false; } - const connection = this.connections.get(serverName); + for (const result of results) { + if (result.status === 'fulfilled') { + const { connection, name } = result.value; - if (!connection) { - console.error(`[MCPStore] No connection found for server: ${serverName}`); + this.connections.set(name, connection); - return null; + this.indexServerTools(name, connection.tools); + } else { + console.error(`[MCPStore] Failed to connect:`, result.reason); + } } - try { - const result = await MCPService.readResource(connection, uri); - const resourceInfo = mcpResourceStore.findResourceByUri(uri); - - if (resourceInfo) { - mcpResourceStore.cacheResourceContent(resourceInfo, result.contents); - } + const successCount = this.connections.size; - return result.contents; - } catch (error) { - console.error(`[MCPStore] Failed to read resource ${uri}:`, error); + if (successCount === 0 && serverEntries.length > 0) { + this.updateState({ + connectedServers: [], + error: 'All MCP server connections failed', + isInitializing: false, + toolCount: 0 + }); + this.initPromise = null; - return null; + return false; } + + this.updateState({ + connectedServers: Array.from(this.connections.keys()), + error: null, + isInitializing: false, + toolCount: this.toolsIndex.size + }); + this.initPromise = null; + + return true; } /** - * Subscribe to resource updates. + * Generates a unique server ID from an optional ID string or index. */ - async subscribeToResource(uri: string): Promise { - const serverName = mcpResourceStore.findServerForUri(uri); - - if (!serverName) { - console.error(`[MCPStore] No server found for resource URI: ${uri}`); - - return false; + private generateServerId(id: unknown, index: number): string { + if (typeof id === 'string' && id.trim()) { + return id.trim(); } + return `${MCP_SERVER_ID_PREFIX}-${index + 1}`; + } + + private handleToolsListChanged(serverName: string, tools: Tool[]): void { const connection = this.connections.get(serverName); if (!connection) { - console.error(`[MCPStore] No connection found for server: ${serverName}`); - - return false; + return; } - if (!MCPService.supportsResourceSubscriptions(connection)) { - return false; + for (const [toolName, ownerServer] of this.toolsIndex.entries()) { + if (ownerServer === serverName) this.toolsIndex.delete(toolName); } - try { - await MCPService.subscribeResource(connection, uri); - mcpResourceStore.addSubscription(uri, serverName); + connection.tools = tools; - return true; - } catch (error) { - console.error(`[MCPStore] Failed to subscribe to resource ${uri}:`, error); + for (const tool of tools) { + if (this.toolsIndex.has(tool.name)) + console.warn( + `[MCPStore] Tool name conflict after list change: "${tool.name}" exists in "${this.toolsIndex.get(tool.name)}" and "${serverName}". Using tool from "${serverName}".` + ); - return false; + this.toolsIndex.set(tool.name, serverName); } + this.updateState({ toolCount: this.toolsIndex.size }); } /** - * Unsubscribe from resource updates. + * Registers the tools exposed by a server into the global name->server index, + * warning on conflicts. Shared by connect, reconnect and auto-reconnect. */ - async unsubscribeFromResource(uri: string): Promise { - const serverName = mcpResourceStore.findServerForUri(uri); - - if (!serverName) { - console.error(`[MCPStore] No server found for resource URI: ${uri}`); + private indexServerTools(serverName: string, tools: Tool[]): void { + for (const tool of tools) { + if (this.toolsIndex.has(tool.name)) + console.warn( + `[MCPStore] Tool name conflict: "${tool.name}" exists in "${this.toolsIndex.get(tool.name)}" and "${serverName}". Using tool from "${serverName}".` + ); - return false; + this.toolsIndex.set(tool.name, serverName); } + } - const connection = this.connections.get(serverName); - - if (!connection) { - console.error(`[MCPStore] No connection found for server: ${serverName}`); - - return false; - } + private async initialize(signature: string, mcpConfig: MCPClientConfig): Promise { + this.updateState({ error: null, isInitializing: true }); + this.configSignature = signature; - try { - await MCPService.unsubscribeResource(connection, uri); - mcpResourceStore.removeSubscription(uri); + const serverEntries = Object.entries(mcpConfig.servers); - return true; - } catch (error) { - console.error(`[MCPStore] Failed to unsubscribe from resource ${uri}:`, error); + if (serverEntries.length === 0) { + this.updateState({ connectedServers: [], isInitializing: false, toolCount: 0 }); return false; } - } - /** - * Add a resource as attachment to chat context. - * Automatically fetches content if not cached. - */ - async attachResource(uri: string): Promise { - const resourceInfo = mcpResourceStore.findResourceByUri(uri); + this.initPromise = this.doInitialize(signature, mcpConfig, serverEntries); - if (!resourceInfo) { - console.error(`[MCPStore] Resource not found: ${uri}`); + return this.initPromise; + } - return null; - } + private parseToolArguments(args: string | Record): Record { + if (typeof args === 'string') { + const trimmed = args.trim(); - // Check if already attached - if (mcpResourceStore.isAttached(uri)) { - return null; - } + if (trimmed === '') { + return {}; + } - // Add attachment (initially loading) - const attachment = mcpResourceStore.addAttachment(resourceInfo); + try { + const parsed = JSON.parse(trimmed); - // Fetch content - try { - const content = await this.readResource(uri); + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) + throw new Error( + `Tool arguments must be an object, got ${Array.isArray(parsed) ? 'array' : typeof parsed}` + ); - if (content) { - mcpResourceStore.updateAttachmentContent(attachment.id, content); - } else { - mcpResourceStore.updateAttachmentError(attachment.id, 'Failed to read resource'); + return parsed as Record; + } catch (error) { + throw new Error(`Failed to parse tool arguments as JSON: ${(error as Error).message}`); } - } catch (error) { - const message = error instanceof Error ? error.message : String(error); + } - mcpResourceStore.updateAttachmentError(attachment.id, message); + if (typeof args === 'object' && args !== null && !Array.isArray(args)) { + return args; } - return mcpResourceStore.getAttachment(attachment.id) ?? null; + throw new Error(`Invalid tool arguments type: ${typeof args}`); } /** - * Remove a resource attachment from chat context. + * Immediately reconnect to a server by creating a fresh transport and session. + * Used when a session-expired error (HTTP 404) is detected during tool execution. + * Per MCP spec 2025-11-25: client MUST discard session ID and re-initialize. + * + * Unlike autoReconnect (which uses exponential backoff for connectivity issues), + * this performs a single immediate reconnection attempt since the server is known + * to be reachable (it responded with 404). */ - removeResourceAttachment(attachmentId: string): void { - mcpResourceStore.removeAttachment(attachmentId); - } + private async reconnectServer(serverName: string): Promise { + const serverConfig = this.serverConfigs.get(serverName); - /** - * Clear all resource attachments. - */ - clearResourceAttachments(): void { - mcpResourceStore.clearAttachments(); - } + if (!serverConfig) { + throw new Error(`[MCPStore] No config found for ${serverName}, cannot reconnect`); + } - /** - * Get formatted resource context for chat. - */ - getResourceContextForChat(): string { - return mcpResourceStore.formatAttachmentsForContext(); + // Disconnect stale connection (clears old transport + session ID) + const oldConnection = this.connections.get(serverName); + + if (oldConnection) { + await MCPService.disconnect(oldConnection).catch(console.warn); + this.connections.delete(serverName); + } + + console.log(`[MCPStore][${serverName}] Session expired, reconnecting with fresh session...`); + + const listChangedHandlers = this.createListChangedHandlers(serverName); + const connection = await MCPService.connect( + serverName, + serverConfig, + DEFAULT_MCP_CONFIG.clientInfo, + DEFAULT_MCP_CONFIG.capabilities, + (phase) => { + if (phase === MCPConnectionPhase.DISCONNECTED) { + console.log(`[MCPStore][${serverName}] Connection lost, starting auto-reconnect`); + this.autoReconnect(serverName); + } + }, + listChangedHandlers + ); + + this.connections.set(serverName, connection); + this.indexServerTools(serverName, connection.tools); + + console.log(`[MCPStore][${serverName}] Session recovered successfully`); } - /** - * Convert current resource attachments to DatabaseMessageExtra[] and clear them. - * Called during message send to persist resources with the user message. - */ - consumeResourceAttachmentsAsExtras(): DatabaseMessageExtraMcpResource[] { - const extras = mcpResourceStore.toMessageExtras(); + private updateState(state: { + isInitializing?: boolean; + error?: string | null; + toolCount?: number; + connectedServers?: string[]; + }): void { + if (state.isInitializing !== undefined) { + this._isInitializing = state.isInitializing; + } - if (extras.length > 0) { - mcpResourceStore.clearAttachments(); + if (state.error !== undefined) { + this._error = state.error; } - return extras; + if (state.toolCount !== undefined) { + this._toolCount = state.toolCount; + } + + if (state.connectedServers !== undefined) { + this.connectedServers = state.connectedServers; + } } } diff --git a/tools/ui/src/lib/stores/mcp-resources.svelte.ts b/tools/ui/src/lib/stores/mcp/resources.svelte.ts similarity index 96% rename from tools/ui/src/lib/stores/mcp-resources.svelte.ts rename to tools/ui/src/lib/stores/mcp/resources.svelte.ts index b68def89f593..79ff2c20927a 100644 --- a/tools/ui/src/lib/stores/mcp-resources.svelte.ts +++ b/tools/ui/src/lib/stores/mcp/resources.svelte.ts @@ -38,32 +38,40 @@ function generateAttachmentId(): string { } class MCPResourceStore { - private _serverResources = $state>(new SvelteMap()); - private _cachedResources = $state>(new SvelteMap()); - private _subscriptions = $state>(new SvelteMap()); private _attachments = $state([]); + private _cachedResources = $state>(new SvelteMap()); private _isLoading = $state(false); + private _serverResources = $state>(new SvelteMap()); + private _subscriptions = $state>(new SvelteMap()); - get serverResources(): Map { - return this._serverResources; + get attachmentCount(): number { + return this._attachments.length; } - get cachedResources(): Map { - return this._cachedResources; + get attachments(): MCPResourceAttachment[] { + return this._attachments; } - get subscriptions(): Map { - return this._subscriptions; + get cachedResources(): Map { + return this._cachedResources; } - get attachments(): MCPResourceAttachment[] { - return this._attachments; + get hasAttachments(): boolean { + return this._attachments.length > 0; } get isLoading(): boolean { return this._isLoading; } + get serverResources(): Map { + return this._serverResources; + } + + get subscriptions(): Map { + return this._subscriptions; + } + get totalResourceCount(): number { let count = 0; @@ -84,134 +92,89 @@ class MCPResourceStore { return count; } - get attachmentCount(): number { - return this._attachments.length; - } - - get hasAttachments(): boolean { - return this._attachments.length > 0; - } - /** - * - * - * Server Resources Management - * - * + * Add a resource attachment to the current chat context */ + addAttachment(resource: MCPResourceInfo): MCPResourceAttachment { + const attachment: MCPResourceAttachment = { + id: generateAttachmentId(), + loading: true, + resource + }; + + this._attachments = [...this._attachments, attachment]; + console.log(`[MCPResources] Added attachment: ${resource.uri}`); + + return attachment; + } /** - * Set resources for a server (called after listResources) + * Register a subscription for a resource */ - setServerResources( - serverName: string, - resources: MCPResource[], - templates: MCPResourceTemplate[] - ): void { - this._serverResources.set(serverName, { - error: undefined, - lastFetched: new Date(), - loading: false, - resources, + addSubscription(uri: string, serverName: string): void { + this._subscriptions.set(uri, { serverName, - templates + subscribedAt: new Date(), + uri }); - console.log( - `[MCPResources][${serverName}] Set ${resources.length} resources, ${templates.length} templates` - ); - } - /** - * Set loading state for a server's resources - */ - setServerLoading(serverName: string, loading: boolean): void { - const existing = this._serverResources.get(serverName); + const cached = this._cachedResources.get(uri); - if (existing) { - this._serverResources.set(serverName, { ...existing, loading }); - } else { - this._serverResources.set(serverName, { - error: undefined, - loading, - resources: [], - serverName, - templates: [] - }); + if (cached) { + this._cachedResources.set(uri, { ...cached, subscribed: true }); } + + console.log(`[MCPResources] Added subscription: ${uri}`); } /** - * Set error state for a server's resources + * Cache resource content after reading */ - setServerError(serverName: string, error: string): void { - const existing = this._serverResources.get(serverName); + cacheResourceContent(resource: MCPResourceInfo, content: MCPResourceContent[]): void { + // Enforce cache size limit + if (this._cachedResources.size >= MCP_RESOURCE_CACHE.MAX_ENTRIES) { + const oldestKey = this._cachedResources.keys().next().value; - if (existing) { - this._serverResources.set(serverName, { ...existing, error, loading: false }); - } else { - this._serverResources.set(serverName, { - error, - loading: false, - resources: [], - serverName, - templates: [] - }); + if (oldestKey) { + this._cachedResources.delete(oldestKey); + } } + + this._cachedResources.set(resource.uri, { + content, + fetchedAt: new Date(), + resource, + subscribed: this._subscriptions.has(resource.uri) + }); + console.log(`[MCPResources] Cached content for: ${resource.uri}`); } /** - * Get resources for a specific server + * Clear all state (e.g., on full reset) */ - getServerResources(serverName: string): MCPServerResources | undefined { - return this._serverResources.get(serverName); + clear(): void { + this._serverResources.clear(); + this._cachedResources.clear(); + this._subscriptions.clear(); + this._attachments = []; + this._isLoading = false; + console.log(`[MCPResources] Cleared all state`); } /** - * Get all resources as MCPResourceInfo array (flattened with server names) + * Clear all attachments */ - getAllResourceInfos(): MCPResourceInfo[] { - const result: MCPResourceInfo[] = []; - - for (const [serverName, serverRes] of this._serverResources) { - for (const resource of serverRes.resources) { - result.push({ - annotations: resource.annotations, - description: resource.description, - icons: resource.icons, - mimeType: resource.mimeType, - name: resource.name, - serverName, - title: resource.title, - uri: resource.uri - }); - } - } - - return result; + clearAttachments(): void { + this._attachments = []; + console.log(`[MCPResources] Cleared all attachments`); } /** - * Get all templates as MCPResourceTemplateInfo array (flattened with server names) + * Clear all cached content */ - getAllTemplateInfos(): MCPResourceTemplateInfo[] { - const result: MCPResourceTemplateInfo[] = []; - - for (const [serverName, serverRes] of this._serverResources) { - for (const template of serverRes.templates) { - result.push({ - annotations: template.annotations, - description: template.description, - icons: template.icons, - mimeType: template.mimeType, - name: template.name, - serverName, - title: template.title, - uriTemplate: template.uriTemplate - }); - } - } - - return result; + clearCache(): void { + this._cachedResources.clear(); + console.log(`[MCPResources] Cleared all cached content`); } /** @@ -220,14 +183,12 @@ class MCPResourceStore { clearServerResources(serverName: string): void { this._serverResources.delete(serverName); - // Also clear cached content for this server's resources for (const [uri, cached] of this._cachedResources) { if (cached.resource.serverName === serverName) { this._cachedResources.delete(uri); } } - // Clear subscriptions for this server for (const [uri, sub] of this._subscriptions) { if (sub.serverName === serverName) { this._subscriptions.delete(uri); @@ -238,152 +199,170 @@ class MCPResourceStore { } /** - * - * - * Resource Content Caching - * - * + * Find resource info by URI across all servers */ + findResourceByUri(uri: string): MCPResourceInfo | undefined { + const normalizedUri = normalizeResourceUri(uri); - /** - * Cache resource content after reading - */ - cacheResourceContent(resource: MCPResourceInfo, content: MCPResourceContent[]): void { - // Enforce cache size limit - if (this._cachedResources.size >= MCP_RESOURCE_CACHE.MAX_ENTRIES) { - // Remove oldest entry - const oldestKey = this._cachedResources.keys().next().value; + for (const [serverName, serverRes] of this._serverResources) { + const resource = + serverRes.resources.find((r) => r.uri === uri) ?? + serverRes.resources.find((r) => normalizeResourceUri(r.uri) === normalizedUri); - if (oldestKey) { - this._cachedResources.delete(oldestKey); + if (resource) { + return { + annotations: resource.annotations, + description: resource.description, + icons: resource.icons, + mimeType: resource.mimeType, + name: resource.name, + serverName, + title: resource.title, + uri: resource.uri + }; } } - this._cachedResources.set(resource.uri, { - content, - fetchedAt: new Date(), - resource, - subscribed: this._subscriptions.has(resource.uri) - }); - console.log(`[MCPResources] Cached content for: ${resource.uri}`); + return undefined; } /** - * Get cached content for a resource + * Find server name for a resource URI */ - getCachedContent(uri: string): MCPCachedResource | undefined { - const cached = this._cachedResources.get(uri); - - if (!cached) return undefined; - - // Check if cache is still valid - const age = Date.now() - cached.fetchedAt.getTime(); - - if (age > MCP_RESOURCE_CACHE.TTL_MS && !cached.subscribed) { - // Cache expired and not subscribed, remove it - this._cachedResources.delete(uri); - - return undefined; + findServerForUri(uri: string): string | undefined { + for (const [serverName, serverRes] of this._serverResources) { + if (serverRes.resources.some((r) => r.uri === uri)) { + return serverName; + } } - return cached; + return undefined; } /** - * Invalidate cached content for a resource (e.g., on update notification) + * Get resource content as text for chat context + * Formats content for inclusion in LLM prompts */ - invalidateCache(uri: string): void { - this._cachedResources.delete(uri); - console.log(`[MCPResources] Invalidated cache for: ${uri}`); - } + formatAttachmentsForContext(): string { + if (this._attachments.length === 0) return ''; - /** - * Clear all cached content - */ - clearCache(): void { - this._cachedResources.clear(); - console.log(`[MCPResources] Cleared all cached content`); - } + const parts: string[] = []; - /** - * - * - * Subscriptions - * - * - */ + for (const attachment of this._attachments) { + if (attachment.error) continue; - /** - * Register a subscription for a resource - */ - addSubscription(uri: string, serverName: string): void { - this._subscriptions.set(uri, { - serverName, - subscribedAt: new Date(), - uri - }); + if (!attachment.content || attachment.content.length === 0) continue; - // Update cached resource if exists - const cached = this._cachedResources.get(uri); + const resourceName = attachment.resource.title || attachment.resource.name; + const serverName = attachment.resource.serverName; - if (cached) { - this._cachedResources.set(uri, { ...cached, subscribed: true }); + for (const content of attachment.content) { + if ('text' in content && content.text) { + parts.push(`\n\n--- Resource: ${resourceName} (from ${serverName}) ---\n${content.text}`); + } else if ('blob' in content && content.blob) { + // For binary content, just note it exists + parts.push( + `\n\n--- Resource: ${resourceName} (from ${serverName}) ---\n[${BINARY_CONTENT_LABEL}: ${content.mimeType || RESOURCE_UNKNOWN_TYPE}]` + ); + } + } } - console.log(`[MCPResources] Added subscription: ${uri}`); + return parts.join(''); } /** - * Remove a subscription for a resource + * Get all resources as MCPResourceInfo array (flattened with server names) */ - removeSubscription(uri: string): void { - this._subscriptions.delete(uri); + getAllResourceInfos(): MCPResourceInfo[] { + const result: MCPResourceInfo[] = []; - // Update cached resource if exists - const cached = this._cachedResources.get(uri); + for (const [serverName, serverRes] of this._serverResources) { + for (const resource of serverRes.resources) { + result.push({ + annotations: resource.annotations, + description: resource.description, + icons: resource.icons, + mimeType: resource.mimeType, + name: resource.name, + serverName, + title: resource.title, + uri: resource.uri + }); + } + } - if (cached) { - this._cachedResources.set(uri, { ...cached, subscribed: false }); + return result; + } + + /** + * Get all templates as MCPResourceTemplateInfo array (flattened with server names) + */ + getAllTemplateInfos(): MCPResourceTemplateInfo[] { + const result: MCPResourceTemplateInfo[] = []; + + for (const [serverName, serverRes] of this._serverResources) { + for (const template of serverRes.templates) { + result.push({ + annotations: template.annotations, + description: template.description, + icons: template.icons, + mimeType: template.mimeType, + name: template.name, + serverName, + title: template.title, + uriTemplate: template.uriTemplate + }); + } } - console.log(`[MCPResources] Removed subscription: ${uri}`); + return result; } /** - * Check if a resource is subscribed + * Get attachment by ID */ - isSubscribed(uri: string): boolean { - return this._subscriptions.has(uri); + getAttachment(attachmentId: string): MCPResourceAttachment | undefined { + return this._attachments.find((att) => att.id === attachmentId); } /** - * Handle resource update notification + * Get cached content for a resource */ - handleResourceUpdate(uri: string): void { - // Invalidate cache so next read gets fresh content - this.invalidateCache(uri); + getCachedContent(uri: string): MCPCachedResource | undefined { + const cached = this._cachedResources.get(uri); - // Update subscription last update time - const sub = this._subscriptions.get(uri); + if (!cached) return undefined; - if (sub) { - this._subscriptions.set(uri, { ...sub, lastUpdate: new Date() }); + const age = Date.now() - cached.fetchedAt.getTime(); + + if (age > MCP_RESOURCE_CACHE.TTL_MS && !cached.subscribed) { + // Cache expired and not subscribed, remove it + this._cachedResources.delete(uri); + + return undefined; } - console.log(`[MCPResources] Resource updated: ${uri}`); + return cached; + } + + /** + * Get resources for a specific server + */ + getServerResources(serverName: string): MCPServerResources | undefined { + return this._serverResources.get(serverName); } /** * Handle resources list changed notification */ handleResourcesListChanged(serverName: string): void { - // Mark server resources as needing refresh const existing = this._serverResources.get(serverName); if (existing) { this._serverResources.set(serverName, { ...existing, - lastFetched: undefined // Mark as stale + lastFetched: undefined }); } @@ -399,60 +378,27 @@ class MCPResourceStore { */ /** - * Add a resource attachment to the current chat context - */ - addAttachment(resource: MCPResourceInfo): MCPResourceAttachment { - const attachment: MCPResourceAttachment = { - id: generateAttachmentId(), - loading: true, - resource - }; - - this._attachments = [...this._attachments, attachment]; - console.log(`[MCPResources] Added attachment: ${resource.uri}`); - - return attachment; - } - - /** - * Update attachment with fetched content + * Handle resource update notification */ - updateAttachmentContent(attachmentId: string, content: MCPResourceContent[]): void { - this._attachments = this._attachments.map((att) => - att.id === attachmentId ? { ...att, content, error: undefined, loading: false } : att - ); - } + handleResourceUpdate(uri: string): void { + // Invalidate cache so next read gets fresh content + this.invalidateCache(uri); - /** - * Update attachment with error - */ - updateAttachmentError(attachmentId: string, error: string): void { - this._attachments = this._attachments.map((att) => - att.id === attachmentId ? { ...att, error, loading: false } : att - ); - } + const sub = this._subscriptions.get(uri); - /** - * Remove an attachment - */ - removeAttachment(attachmentId: string): void { - this._attachments = this._attachments.filter((att) => att.id !== attachmentId); - console.log(`[MCPResources] Removed attachment: ${attachmentId}`); - } + if (sub) { + this._subscriptions.set(uri, { ...sub, lastUpdate: new Date() }); + } - /** - * Clear all attachments - */ - clearAttachments(): void { - this._attachments = []; - console.log(`[MCPResources] Cleared all attachments`); + console.log(`[MCPResources] Resource updated: ${uri}`); } /** - * Get attachment by ID + * Invalidate cached content for a resource (e.g., on update notification) */ - getAttachment(attachmentId: string): MCPResourceAttachment | undefined { - return this._attachments.find((att) => att.id === attachmentId); + invalidateCache(uri: string): void { + this._cachedResources.delete(uri); + console.log(`[MCPResources] Invalidated cache for: ${uri}`); } /** @@ -467,103 +413,99 @@ class MCPResourceStore { } /** - * - * - * Utility Methods - * - * + * Check if a resource is subscribed */ + isSubscribed(uri: string): boolean { + return this._subscriptions.has(uri); + } /** - * Set global loading state + * Remove an attachment */ - setLoading(loading: boolean): void { - this._isLoading = loading; + removeAttachment(attachmentId: string): void { + this._attachments = this._attachments.filter((att) => att.id !== attachmentId); + console.log(`[MCPResources] Removed attachment: ${attachmentId}`); } /** - * Find resource info by URI across all servers + * Remove a subscription for a resource */ - findResourceByUri(uri: string): MCPResourceInfo | undefined { - const normalizedUri = normalizeResourceUri(uri); + removeSubscription(uri: string): void { + this._subscriptions.delete(uri); - for (const [serverName, serverRes] of this._serverResources) { - const resource = - serverRes.resources.find((r) => r.uri === uri) ?? - serverRes.resources.find((r) => normalizeResourceUri(r.uri) === normalizedUri); + const cached = this._cachedResources.get(uri); - if (resource) { - return { - annotations: resource.annotations, - description: resource.description, - icons: resource.icons, - mimeType: resource.mimeType, - name: resource.name, - serverName, - title: resource.title, - uri: resource.uri - }; - } + if (cached) { + this._cachedResources.set(uri, { ...cached, subscribed: false }); } - return undefined; + console.log(`[MCPResources] Removed subscription: ${uri}`); } /** - * Find server name for a resource URI + * Set global loading state */ - findServerForUri(uri: string): string | undefined { - for (const [serverName, serverRes] of this._serverResources) { - if (serverRes.resources.some((r) => r.uri === uri)) { - return serverName; - } - } - - return undefined; + setLoading(loading: boolean): void { + this._isLoading = loading; } /** - * Clear all state (e.g., on full reset) + * Set error state for a server's resources */ - clear(): void { - this._serverResources.clear(); - this._cachedResources.clear(); - this._subscriptions.clear(); - this._attachments = []; - this._isLoading = false; - console.log(`[MCPResources] Cleared all state`); + setServerError(serverName: string, error: string): void { + const existing = this._serverResources.get(serverName); + + if (existing) { + this._serverResources.set(serverName, { ...existing, error, loading: false }); + } else { + this._serverResources.set(serverName, { + error, + loading: false, + resources: [], + serverName, + templates: [] + }); + } } /** - * Get resource content as text for chat context - * Formats content for inclusion in LLM prompts + * Set loading state for a server's resources */ - formatAttachmentsForContext(): string { - if (this._attachments.length === 0) return ''; - - const parts: string[] = []; - - for (const attachment of this._attachments) { - if (attachment.error) continue; - - if (!attachment.content || attachment.content.length === 0) continue; - - const resourceName = attachment.resource.title || attachment.resource.name; - const serverName = attachment.resource.serverName; + setServerLoading(serverName: string, loading: boolean): void { + const existing = this._serverResources.get(serverName); - for (const content of attachment.content) { - if ('text' in content && content.text) { - parts.push(`\n\n--- Resource: ${resourceName} (from ${serverName}) ---\n${content.text}`); - } else if ('blob' in content && content.blob) { - // For binary content, just note it exists - parts.push( - `\n\n--- Resource: ${resourceName} (from ${serverName}) ---\n[${BINARY_CONTENT_LABEL}: ${content.mimeType || RESOURCE_UNKNOWN_TYPE}]` - ); - } - } + if (existing) { + this._serverResources.set(serverName, { ...existing, loading }); + } else { + this._serverResources.set(serverName, { + error: undefined, + loading, + resources: [], + serverName, + templates: [] + }); } + } - return parts.join(''); + /** + * Set resources for a server (called after listResources) + */ + setServerResources( + serverName: string, + resources: MCPResource[], + templates: MCPResourceTemplate[] + ): void { + this._serverResources.set(serverName, { + error: undefined, + lastFetched: new Date(), + loading: false, + resources, + serverName, + templates + }); + console.log( + `[MCPResources][${serverName}] Set ${resources.length} resources, ${templates.length} templates` + ); } /** @@ -605,6 +547,24 @@ class MCPResourceStore { return extras; } + + /** + * Update attachment with fetched content + */ + updateAttachmentContent(attachmentId: string, content: MCPResourceContent[]): void { + this._attachments = this._attachments.map((att) => + att.id === attachmentId ? { ...att, content, error: undefined, loading: false } : att + ); + } + + /** + * Update attachment with error + */ + updateAttachmentError(attachmentId: string, error: string): void { + this._attachments = this._attachments.map((att) => + att.id === attachmentId ? { ...att, error, loading: false } : att + ); + } } export const mcpResourceStore = new MCPResourceStore(); diff --git a/tools/ui/src/lib/stores/models.svelte.ts b/tools/ui/src/lib/stores/models.svelte.ts deleted file mode 100644 index c741d144c402..000000000000 --- a/tools/ui/src/lib/stores/models.svelte.ts +++ /dev/null @@ -1,1077 +0,0 @@ -import { FAVORITE_MODELS_LOCALSTORAGE_KEY, MODEL_PROPS_CACHE } from '$lib/constants'; -import { - FileTypeCategory, - ModelModality, - ServerModelsSseEventType, - ServerModelStatus -} from '$lib/enums'; -import { ModelsService } from '$lib/services/models.service'; -import { PropsService } from '$lib/services/props.service'; -// direct imports between stores, not via the barrel, to avoid circular deps -import { conversationsStore } from '$lib/stores/conversations.svelte'; -import { serverStore } from '$lib/stores/server.svelte'; -// deep imports, not the '$lib/utils' barrel: it re-exports modules that reach back -// into the stores, and going through it here would read a half-built module -import { TTLCache } from '$lib/utils/cache-ttl'; -import { - detectThinkingSupport, - detectThinkingSupportWithReason -} from '$lib/utils/chat-template-thinking-detector'; -import { getConversationModel } from '$lib/utils/conversation-utils'; -import { SvelteMap, SvelteSet } from 'svelte/reactivity'; -import { toast } from 'svelte-sonner'; - -/** - * modelsStore - Reactive store for model management in both MODEL and ROUTER modes. - * - * **Architecture & Relationships:** - * - **ModelsService**: Stateless service for model API communication - * - **PropsService**: Stateless service for props/modalities fetching - * - **modelsStore** (this class): Reactive store for model state - * - **conversationsStore**: Tracks which conversations use which models - * - * **API Inconsistency Workaround:** - * In MODEL mode, `/props` returns modalities for the single model. - * In ROUTER mode, `/props` has no modalities — must use `/props?model=` per model. - * This store normalizes this behavior so consumers don't need to know the server mode. - */ -class ModelsStore { - /** - * - * - * State - * - * - */ - - models = $state([]); - routerModels = $state([]); - loading = $state(false); - updating = $state(false); - error = $state(null); - selectedModelId = $state(null); - selectedModelName = $state(null); - - // Dedup concurrent fetch() callers — all awaiters share the same inflight promise. - // Without this, ?model= URL handler races an in-progress fetch and sees an empty list. - private inflightFetch: Promise | null = null; - - private modelUsage = $state>>(new Map()); - private modelLoadingStates = new SvelteMap(); - - // /models/sse feed state, the single source of truth for status and load progress - private statusAbort: AbortController | null = null; - private statusReaderActive = false; - private loadProgress = new SvelteMap(); - private statusWaiters = new Map< - string, - { target: ServerModelStatus; resolve: () => void; reject: (e: Error) => void } - >(); - - favoriteModelIds = $state>(this.loadFavoritesFromStorage()); - - /** - * Model-specific props cache with TTL. - * Key: modelId, Value: props data including modalities. - * TTL: 10 minutes — props don't change frequently. - */ - private modelPropsCache = new TTLCache({ - maxEntries: MODEL_PROPS_CACHE.MAX_ENTRIES, - ttlMs: MODEL_PROPS_CACHE.TTL_MS - }); - private modelPropsFetching = $state>(new Set()); - - /** - * Version counter for props cache — used to trigger reactivity when props are updated. - */ - propsCacheVersion = $state(0); - - /** - * - * - * Computed Getters - * - * - */ - - get selectedModel(): ModelOption | null { - if (!this.selectedModelId) return null; - - return this.models.find((m) => m.id === this.selectedModelId) ?? null; - } - - get loadedModelIds(): string[] { - return this.routerModels - .filter( - (m) => - m.status.value === ServerModelStatus.LOADED || - m.status.value === ServerModelStatus.SLEEPING - ) - .map((m) => m.id); - } - - get loadingModelIds(): string[] { - return Array.from(this.modelLoadingStates.entries()) - .filter(([, loading]) => loading) - .map(([id]) => id); - } - - /** - * Get model name in MODEL mode (single model). - * Extracts from model_path or model_alias from server props. - * In ROUTER mode, returns null (model is per-conversation). - */ - get singleModelName(): string | null { - if (serverStore.isRouterMode) return null; - - const props = serverStore.props; - - if (props?.model_alias) return props.model_alias; - - if (!props?.model_path) return null; - - return props.model_path.split(/(\\|\/)/).pop() || null; - } - - /** - * Model the active conversation view resolves to. Router mode: the user's - * selection first, then the conversation's own model. Otherwise the single - * served model, from the models list or the server props as a fallback. - */ - get activeModelId(): string | null { - if (!serverStore.isRouterMode) { - return this.models.length > 0 ? this.models[0].model : this.singleModelName; - } - - if (this.selectedModelId) { - const selected = this.models.find((m) => m.id === this.selectedModelId); - - if (selected) return selected.model; - } - - const conversationModel = getConversationModel(conversationsStore.activeMessages); - - if (conversationModel) { - const model = this.models.find((m) => m.model === conversationModel); - - if (model) return model.model; - } - - return null; - } - - get selectedModelContextSize(): number | null { - if (!this.selectedModelName) return null; - - return this.getModelContextSize(this.selectedModelName); - } - - /** - * - * - * Modalities - * - * - */ - - getModelModalities(modelId: string): ModelModalities | null { - if (!serverStore.isRouterMode && serverStore.props?.modalities) { - return this.buildModalities(serverStore.props.modalities); - } - - const model = this.models.find((m) => m.model === modelId || m.id === modelId); - - if (model?.modalities) { - return model.modalities; - } - - const props = this.modelPropsCache.get(modelId); - - if (props?.modalities) { - return this.buildModalities(props.modalities); - } - - return null; - } - - modelSupportsVision(modelId: string): boolean { - return this.getModelModalities(modelId)?.vision ?? false; - } - - modelSupportsAudio(modelId: string): boolean { - return this.getModelModalities(modelId)?.audio ?? false; - } - - modelSupportsVideo(modelId: string): boolean { - return this.getModelModalities(modelId)?.video ?? false; - } - - getModelModalitiesArray(modelId: string): ModelModality[] { - const modalities = this.getModelModalities(modelId); - - if (!modalities) return []; - - const result: ModelModality[] = []; - - if (modalities.vision) result.push(ModelModality.VISION); - - if (modalities.audio) result.push(ModelModality.AUDIO); - - if (modalities.video) result.push(ModelModality.VIDEO); - - return result; - } - - getModelProps(modelId: string): ApiLlamaCppServerProps | null { - return this.modelPropsCache.get(modelId); - } - - getModelContextSize(modelId: string): number | null { - const props = this.getModelProps(modelId); - const nCtx = props?.default_generation_settings?.n_ctx; - - return typeof nCtx === 'number' ? nCtx : null; - } - - isModelPropsFetching(modelId: string): boolean { - return this.modelPropsFetching.has(modelId); - } - - /** - * - * - * Status Queries - * - * - */ - - isModelLoaded(modelId: string): boolean { - const model = this.routerModels.find((m) => m.id === modelId); - - return ( - model?.status.value === ServerModelStatus.LOADED || - model?.status.value === ServerModelStatus.SLEEPING - ); - } - - isModelOperationInProgress(modelId: string): boolean { - return this.modelLoadingStates.get(modelId) ?? false; - } - - getModelStatus(modelId: string): ServerModelStatus | null { - const model = this.routerModels.find((m) => m.id === modelId); - - return model?.status.value ?? null; - } - - getModelUsage(modelId: string): SvelteSet { - return this.modelUsage.get(modelId) ?? new SvelteSet(); - } - - isModelInUse(modelId: string): boolean { - const usage = this.modelUsage.get(modelId); - - return usage !== undefined && usage.size > 0; - } - // - // Thinking Support Detection - // - - /** - * Whether the selected model's chat template supports thinking/reasoning. - * Uses heuristic detection on the model's chat_template from /props. - * - * - MODEL mode: the global /props already describes the single loaded model, - * so its chat_template is used directly and no per-model cache is involved - * - ROUTER mode: fetches /props?model= for the selected model (cached), - * triggering an async fetch if not yet cached - */ - get supportsThinking(): boolean { - if (!serverStore.isRouterMode) { - return detectThinkingSupport(serverStore.props?.chat_template ?? ''); - } - - const modelId = this.selectedModelName; - - if (!modelId) return false; - - if (!this.modelPropsCache.get(modelId)) { - this.fetchModelProps(modelId); - } - - const props = this.getModelProps(modelId); - - return detectThinkingSupport(props?.chat_template ?? ''); - } - - /** - * Check if a specific model supports thinking. - * In MODEL mode the global /props describes the single loaded model. - * In ROUTER mode, fetches model props if not cached. - */ - checkModelSupportsThinking(modelId: string): boolean { - if (!serverStore.isRouterMode) { - return detectThinkingSupport(serverStore.props?.chat_template ?? ''); - } - - if (!modelId) return false; - - if (!this.modelPropsCache.get(modelId)) { - this.fetchModelProps(modelId); - } - - const props = this.getModelProps(modelId); - - return detectThinkingSupport(props?.chat_template ?? ''); - } - - /** - * Detailed thinking support detection result with reason for debugging/UI. - */ - get thinkingSupportDetails(): { supported: boolean; reason: string } { - if (!serverStore.isRouterMode) { - return detectThinkingSupportWithReason(serverStore.props?.chat_template ?? ''); - } - - const modelId = this.selectedModelName; - - if (!modelId) { - return { reason: 'No model selected', supported: false }; - } - - if (!this.modelPropsCache.get(modelId)) { - this.fetchModelProps(modelId); - } - - const props = this.getModelProps(modelId); - - return detectThinkingSupportWithReason(props?.chat_template ?? ''); - } - - /** - * - * - * Data Fetching - * - * - */ - - /** - * Fetch list of models from server and detect server role. - * Also fetches modalities for MODEL mode (single model). - */ - async fetch(force = false): Promise { - if (this.inflightFetch) return this.inflightFetch; - - if (this.models.length > 0 && !force) return; - - this.inflightFetch = this.runFetch(); - try { - await this.inflightFetch; - } finally { - this.inflightFetch = null; - } - } - - private async runFetch(): Promise { - this.loading = true; - this.error = null; - - try { - if (!serverStore.props) { - await serverStore.fetch(); - } - - const router = serverStore.isRouterMode; - - if (router) { - const response = await ModelsService.listRouter(); - - this.routerModels = response.data; - this.models = this.buildModelOptions(response); - - await this.fetchModalitiesForLoadedModels(); - - const visible = this.getVisibleModels(); - - if (visible.length === 1 && this.isModelLoaded(visible[0].model)) { - this.selectModelById(visible[0].id); - } - } else { - this.models = await this.fetchModelModeInternal(); - } - } catch (error) { - this.models = []; - this.error = error instanceof Error ? error.message : 'Failed to load models'; - - throw error; - } finally { - this.loading = false; - } - } - - /** Fetch models in MODEL mode (single model, standard OpenAI-compatible). */ - private async fetchModelModeInternal(): Promise { - const response = await ModelsService.list(); - - return this.buildModelOptions(response); - } - - /** - * Build ModelOption[] from an API response. - * Both MODEL and ROUTER modes share the same mapping logic; - * they differ only in which endpoint is called. - */ - private buildModelOptions( - response: ApiModelListResponse | ApiRouterModelsListResponse - ): ModelOption[] { - return response.data.map((item: ApiModelDataEntry, index: number) => { - const details = response.models?.[index]; - const rawCapabilities = Array.isArray(details?.capabilities) ? details?.capabilities : []; - const displayNameSource = - details?.name && details.name.trim().length > 0 ? details.name : item.id; - const modelId = details?.model || item.id; - - return { - aliases: item.aliases ?? [], - capabilities: rawCapabilities.filter((value: unknown): value is string => Boolean(value)), - description: details?.description, - details: details?.details, - id: item.id, - meta: item.meta ?? null, - modalities: this.buildArchitectureModalities(item.architecture), - model: modelId, - name: this.toDisplayName(displayNameSource), - parsedId: ModelsService.parseModelId(modelId), - tags: item.tags ?? [] - }; - }); - } - - /** - * Fetch router models with full metadata (ROUTER mode only). - * No-op in router mode — fetch() already calls listRouter() internally. - * Kept for API compatibility (e.g. handleOpenChange dropdown open handler). - */ - async fetchRouterModels(): Promise { - if (!serverStore.isRouterMode) return; - - try { - const response = await ModelsService.listRouter(); - - this.routerModels = response.data; - await this.fetchModalitiesForLoadedModels(); - - const visible = this.getVisibleModels(); - - if (visible.length === 1 && this.isModelLoaded(visible[0].model)) { - this.selectModelById(visible[0].id); - } - } catch (error) { - console.warn('Failed to fetch router models:', error); - this.routerModels = []; - } - } - - /** - * Fetch props for a specific model from /props endpoint. - * Uses caching to avoid redundant requests. - * - * In ROUTER mode, this only fetches props if the model is loaded, - * since unloaded models return 400 from /props endpoint. - * - * @param modelId - Model identifier to fetch props for - * @returns Props data or null if fetch failed or model not loaded - */ - async fetchModelProps(modelId: string): Promise { - const cached = this.modelPropsCache.get(modelId); - - if (cached) return cached; - - if (serverStore.isRouterMode && !this.isModelLoaded(modelId)) { - return null; - } - - if (this.modelPropsFetching.has(modelId)) return null; - - this.modelPropsFetching.add(modelId); - - try { - const props = await PropsService.fetchForModel(modelId); - - this.modelPropsCache.set(modelId, props); - this.propsCacheVersion++; - - return props; - } catch (error) { - console.warn(`Failed to fetch props for model ${modelId}:`, error); - - return null; - } finally { - this.modelPropsFetching.delete(modelId); - } - } - - /** Fetch modalities for all loaded models from /props endpoint. */ - async fetchModalitiesForLoadedModels(): Promise { - const loadedModelIds = this.loadedModelIds; - - if (loadedModelIds.length === 0) return; - - const propsPromises = loadedModelIds.map((modelId) => this.fetchModelProps(modelId)); - - try { - const results = await Promise.all(propsPromises); - - this.models = this.models.map((model) => { - const modelIndex = loadedModelIds.indexOf(model.model); - - if (modelIndex === -1) return model; - - const props = results[modelIndex]; - - if (!props?.modalities) return model; - - return { ...model, modalities: this.buildModalities(props.modalities) }; - }); - - this.propsCacheVersion++; - } catch (error) { - console.warn('Failed to fetch modalities for loaded models:', error); - } - } - - /** - * Update modalities for a specific model. - * Called when a model is loaded or when we need fresh modality data. - */ - async updateModelModalities(modelId: string): Promise { - const props = await this.fetchModelProps(modelId); - - if (!props?.modalities) return; - - this.models = this.models.map((model) => - model.model === modelId - ? { ...model, modalities: this.buildModalities(props.modalities!) } - : model - ); - - this.propsCacheVersion++; - } - - /** - * Filter to models visible in the UI (ui !== false). - */ - private getVisibleModels(): ModelOption[] { - return this.models.filter((option) => this.getModelProps(option.model)?.ui !== false); - } - - /** - * Gets the model name from the last assistant message in the active conversation. - * Used by both the chat page and settings page to maintain model consistency. - */ - getModelFromLastAssistantResponse(): string | null { - const messages = conversationsStore.activeMessages; - - if (!messages || messages.length === 0) return null; - - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].model) { - return messages[i].model; - } - } - - return null; - } - - /** - * Auto-selects the model from the last assistant response if available and loaded. - * Returns true if a model was selected, false otherwise. - */ - async selectModelFromLastAssistantResponse(): Promise { - const lastModel = this.getModelFromLastAssistantResponse(); - - if (!lastModel || this.selectedModelName === lastModel) return false; - - const matchingModel = this.models.find((option) => option.model === lastModel); - - if (!matchingModel || !this.isModelLoaded(lastModel)) return false; - - try { - await this.selectModelById(matchingModel.id); - console.log(`[modelsStore] Automatically selected model: ${lastModel} from last message`); - - return true; - } catch (error) { - console.warn('[modelsStore] Failed to automatically select model from last message:', error); - - return false; - } - } - - /** - * Auto-selects the first available model if none is selected. - * Prioritizes: - * 1. Model from active conversation's last assistant response (if loaded) - * 2. Model from active conversation's last assistant response (if not loaded) - * 3. First loaded model (not from active conversation) - * 4. A favorite model - * 5. First available model - */ - async ensureFirstModelSelected(): Promise { - if (this.selectedModelName) return; - - const availableModels = this.getVisibleModels(); - - if (availableModels.length === 0) return; - - // Try to select model from last assistant response first - const lastModel = this.getModelFromLastAssistantResponse(); - - if (lastModel) { - const lastModelOption = availableModels.find((m) => m.model === lastModel); - - if (lastModelOption) { - await this.selectModelById(lastModelOption.id); - - if (this.isModelLoaded(lastModel)) { - await this.fetchModelProps(lastModel); - } - - return; - } - } - - // Try a loaded model first - const loadedModel = availableModels.find((m) => this.isModelLoaded(m.model)); - - if (loadedModel) { - await this.selectModelById(loadedModel.id); - await this.fetchModelProps(loadedModel.model); - - return; - } - - // Try loading a favorite model - const favorite = this.favoriteModelIds.values().next()?.value; - - if (favorite) { - await this.selectModelById(favorite); - - return; - } - - // Fall back to the first available model - await this.selectModelById(availableModels[0].id); - } - - /** - * - * - * Model Selection - * - * - */ - - async selectModelById(modelId: string): Promise { - if (!modelId || this.updating) return; - - if (this.selectedModelId === modelId) return; - - const option = this.models.find((model) => model.id === modelId); - - if (!option) throw new Error('Selected model is not available'); - - this.updating = true; - this.error = null; - - try { - this.selectedModelId = option.id; - this.selectedModelName = option.model; - } finally { - this.updating = false; - } - } - - /** - * Select a model by its model name (used for syncing with conversation model). - */ - selectModelByName(modelName: string): void { - const option = this.models.find((model) => model.model === modelName); - - if (option) { - this.selectedModelId = option.id; - this.selectedModelName = option.model; - } - } - - clearSelection(): void { - this.selectedModelId = null; - this.selectedModelName = null; - } - - findModelByName(modelName: string): ModelOption | null { - return ( - this.models.find( - (model) => - model.model === modelName || model.id === modelName || model.aliases?.includes(modelName) - ) ?? null - ); - } - - findModelById(modelId: string): ModelOption | null { - return this.models.find((model) => model.id === modelId) ?? null; - } - - hasModel(modelName: string): boolean { - return this.models.some((model) => model.model === modelName); - } - - /** - * - * - * Loading / Unloading Models - * - * - */ - - // reconnect delay after the feed drops or the server is not ready yet - /** - * Open the /models/sse feed and keep it live with auto reconnect. - * Idempotent and router mode only. The feed drives status and progress, - * so it replaces any post-operation polling. - */ - subscribeStatus(): void { - if (this.statusReaderActive) return; - - if (!serverStore.isRouterMode) return; - - this.statusReaderActive = true; - this.statusAbort = new AbortController(); - void this.runStatusReader(this.statusAbort.signal); - } - - /** - * Close the /models/sse feed and drop transient progress. - */ - unsubscribeStatus(): void { - this.statusReaderActive = false; - this.statusAbort?.abort(); - this.statusAbort = null; - this.loadProgress.clear(); - } - - /** - * Current load progress for a model, or null when not loading. - */ - getLoadProgress(modelId: string): ModelLoadProgress | null { - return this.loadProgress.get(modelId) ?? null; - } - - /** - * Read the feed and reconnect until unsubscribed. - */ - private async runStatusReader(signal: AbortSignal): Promise { - await ModelsService.watchModelEvents(signal, (event) => this.applyStatusEvent(event)); - } - - /** - * Route one feed record by event kind. Only the status_* events carry a - * status payload, models_reload triggers a list refresh, model_remove drops - * the row, download_* belong to the download surface, not here. - */ - private applyStatusEvent(event: ApiModelsSseEvent): void { - switch (event.event) { - case ServerModelsSseEventType.STATUS_CHANGE: - case ServerModelsSseEventType.MODEL_STATUS: - case ServerModelsSseEventType.STATUS_UPDATE: - this.applyModelStatus(event); - - break; - case ServerModelsSseEventType.MODELS_RELOAD: - void this.fetchRouterModels(); - - break; - case ServerModelsSseEventType.MODEL_REMOVE: - this.removeRouterModel(event.model); - - break; - case ServerModelsSseEventType.DOWNLOAD_PROGRESS: - break; - } - } - - /** - * Apply a status envelope: update the model row, track or clear progress, - * settle any pending load or unload awaiter. - */ - private applyModelStatus(event: ApiModelsSseEvent): void { - const model = event.model; - const data = event.data; - - if (!model || !data?.status) return; - - const status = data.status; - - this.setRouterModelStatus(model, status); - - if (status === ServerModelStatus.LOADING) { - if (data.progress) this.loadProgress.set(model, data.progress); - } else { - this.loadProgress.delete(model); - } - - if (status === ServerModelStatus.LOADED) { - void this.updateModelModalities(model); - } - - const failed = - status === ServerModelStatus.FAILED || - (status === ServerModelStatus.UNLOADED && (data.exit_code ?? 0) !== 0); - - if (failed) { - this.rejectStatus(model, new Error(`Model failed: ${this.toDisplayName(model)}`)); - - return; - } - - this.settleStatus(model, status); - } - - /** - * Drop a model row reported gone by the feed and settle its awaiters. - */ - private removeRouterModel(modelId: string): void { - if (this.routerModels.findIndex((m) => m.id === modelId) === -1) return; - - this.routerModels = this.routerModels.filter((m) => m.id !== modelId); - this.loadProgress.delete(modelId); - this.rejectStatus(modelId, new Error(`Model removed: ${this.toDisplayName(modelId)}`)); - } - - /** - * Update one model row status in place, reassigning to trigger reactivity. - */ - private setRouterModelStatus(modelId: string, status: ServerModelStatus): void { - const idx = this.routerModels.findIndex((m) => m.id === modelId); - - if (idx === -1) return; - - const current = this.routerModels[idx]; - - if (current.status.value === status) return; - - const next = [...this.routerModels]; - - next[idx] = { ...current, status: { ...current.status, value: status } }; - this.routerModels = next; - } - - /** - * Register an awaiter that resolves when the feed reports target status. - * One operation runs per model at a time, so one awaiter per model is kept. - */ - private waitForStatus(modelId: string, target: ServerModelStatus): Promise { - return new Promise((resolve, reject) => { - this.statusWaiters.set(modelId, { reject, resolve, target }); - }); - } - - /** - * Resolve and drop the awaiter when the model reaches its target status. - */ - private settleStatus(modelId: string, status: ServerModelStatus): void { - const waiter = this.statusWaiters.get(modelId); - - if (waiter && waiter.target === status) { - this.statusWaiters.delete(modelId); - waiter.resolve(); - } - } - - /** - * Reject and drop the awaiter for a model. - */ - private rejectStatus(modelId: string, error: Error): void { - const waiter = this.statusWaiters.get(modelId); - - if (waiter) { - this.statusWaiters.delete(modelId); - waiter.reject(error); - } - } - - async loadModel(modelId: string): Promise { - if (this.isModelLoaded(modelId)) return; - - if (this.modelLoadingStates.get(modelId)) return; - - this.modelLoadingStates.set(modelId, true); - this.error = null; - - // the feed drives completion, so it must be live before the request - this.subscribeStatus(); - - const reachedLoaded = this.waitForStatus(modelId, ServerModelStatus.LOADED); - - reachedLoaded.catch(() => {}); - - try { - await ModelsService.load(modelId); - await reachedLoaded; - toast.success(`Model loaded: ${this.toDisplayName(modelId)}`); - } catch (error) { - this.rejectStatus(modelId, error instanceof Error ? error : new Error('load failed')); - this.error = error instanceof Error ? error.message : 'Failed to load model'; - toast.error(`Failed to load model: ${this.toDisplayName(modelId)}`); - - throw error; - } finally { - this.modelLoadingStates.set(modelId, false); - } - } - - async unloadModel(modelId: string): Promise { - if (!this.isModelLoaded(modelId)) return; - - if (this.modelLoadingStates.get(modelId)) return; - - this.modelLoadingStates.set(modelId, true); - this.error = null; - - this.subscribeStatus(); - - const reachedUnloaded = this.waitForStatus(modelId, ServerModelStatus.UNLOADED); - - reachedUnloaded.catch(() => {}); - - try { - await ModelsService.unload(modelId); - await reachedUnloaded; - toast.info(`Model unloaded: ${this.toDisplayName(modelId)}`); - } catch (error) { - this.rejectStatus(modelId, error instanceof Error ? error : new Error('unload failed')); - this.error = error instanceof Error ? error.message : 'Failed to unload model'; - toast.error(`Failed to unload model: ${this.toDisplayName(modelId)}`); - - throw error; - } finally { - this.modelLoadingStates.set(modelId, false); - } - } - - async ensureModelLoaded(modelId: string): Promise { - if (this.isModelLoaded(modelId)) return; - - await this.loadModel(modelId); - } - - /** - * - * - * Favorites - * - * - */ - - isFavorite(modelId: string): boolean { - return this.favoriteModelIds.has(modelId); - } - - toggleFavorite(modelId: string): void { - const next = new SvelteSet(this.favoriteModelIds); - - if (next.has(modelId)) { - next.delete(modelId); - } else { - next.add(modelId); - } - - this.favoriteModelIds = next; - - try { - localStorage.setItem(FAVORITE_MODELS_LOCALSTORAGE_KEY, JSON.stringify([...next])); - } catch { - toast.error('Failed to save favorite models to local storage'); - } - } - - private loadFavoritesFromStorage(): Set { - try { - const raw = localStorage.getItem(FAVORITE_MODELS_LOCALSTORAGE_KEY); - - return raw ? new Set(JSON.parse(raw) as string[]) : new Set(); - } catch { - toast.error('Failed to load favorite models from local storage'); - - return new Set(); - } - } - - /** - * - * - * Utilities - * - * - */ - - private toDisplayName(id: string): string { - const segments = id.split(/\\|\//); - const candidate = segments.pop(); - - return candidate && candidate.trim().length > 0 ? candidate : id; - } - - private buildModalities( - modalities: NonNullable - ): ModelModalities { - return { - audio: modalities.audio ?? false, - video: modalities.video ?? false, - vision: modalities.vision ?? false - }; - } - - /** Map the router modalities, the only source available while a model is not loaded. */ - private buildArchitectureModalities( - architecture: ApiModelDataEntry['architecture'] - ): ModelModalities | undefined { - if (!architecture) return undefined; - - const inputs = architecture.input_modalities; - - return { - audio: inputs.includes(FileTypeCategory.AUDIO), - video: inputs.includes(FileTypeCategory.VIDEO), - vision: inputs.includes(FileTypeCategory.IMAGE) - }; - } - - clear(): void { - this.unsubscribeStatus(); - this.statusWaiters.forEach((waiter) => waiter.reject(new Error('Models store cleared'))); - this.statusWaiters.clear(); - this.models = []; - this.routerModels = []; - this.loading = false; - this.updating = false; - this.error = null; - this.selectedModelId = null; - this.selectedModelName = null; - this.modelUsage.clear(); - this.modelLoadingStates.clear(); - this.modelPropsCache.clear(); - this.modelPropsFetching.clear(); - } - - /** - * Prune expired entries from caches. - * Call periodically for proactive memory cleanup. - */ - pruneExpiredCache(): number { - return this.modelPropsCache.prune(); - } -} - -export const modelsStore = new ModelsStore(); diff --git a/tools/ui/src/lib/stores/models/index.svelte.ts b/tools/ui/src/lib/stores/models/index.svelte.ts new file mode 100644 index 000000000000..90d6fe76b719 --- /dev/null +++ b/tools/ui/src/lib/stores/models/index.svelte.ts @@ -0,0 +1,451 @@ +/** + * modelsStore - Model management for MODEL and ROUTER modes + * + * Owns model lists, selection, favorites and load/unload state. Composes the + * per-model props cache (modalities, thinking detection) as + * {@link ModelsStore.props} and the /models/sse status feed as + * {@link ModelsStore.status}; tracks which conversations use which models. + */ + +import { FAVORITE_MODELS_LOCALSTORAGE_KEY } from '$lib/constants'; +import { ServerModelStatus } from '$lib/enums'; +import { ModelsService } from '$lib/services/models.service'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; +import { type ModelPropsHost, ModelPropsManager } from '$lib/stores/models/props.svelte'; +import { type ModelStatusHost, ModelStatusManager } from '$lib/stores/models/status.svelte'; +import { serverStore } from '$lib/stores/server.svelte'; +import { getConversationModel } from '$lib/utils/conversation-utils'; +import { SvelteSet } from 'svelte/reactivity'; +import { toast } from 'svelte-sonner'; + +class ModelsStore implements ModelPropsHost, ModelStatusHost { + error = $state(null); + favoriteModelIds = $state>(this.loadFavoritesFromStorage()); + loading = $state(false); + models = $state([]); + routerModels = $state([]); + selectedModelId = $state(null); + selectedModelName = $state(null); + + updating = $state(false); + + /** Per-model props cache, modalities and thinking detection, composed here. */ + private _props = new ModelPropsManager(this); + + /** Load/unload operations and the /models/sse status feed, composed here. */ + private _status = new ModelStatusManager(this); + + // Dedup concurrent fetch() callers — all awaiters share the same inflight promise. + // Without this, ?model= URL handler races an in-progress fetch and sees an empty list. + private inflightFetch: Promise | null = null; + + /** + * Model the active conversation view resolves to. Router mode: the user's + * selection first, then the conversation's own model. Otherwise the single + * served model, from the models list or the server props as a fallback. + */ + get activeModelId(): string | null { + if (!serverStore.isRouterMode) { + return this.models.length > 0 ? this.models[0].model : this.singleModelName; + } + + if (this.selectedModelId) { + const selected = this.models.find((m) => m.id === this.selectedModelId); + + if (selected) return selected.model; + } + + const conversationModel = getConversationModel(conversationsStore.activeMessages); + + if (conversationModel) { + const model = this.models.find((m) => m.model === conversationModel); + + if (model) return model.model; + } + + return null; + } + + get loadedModelIds(): string[] { + return this.routerModels + .filter( + (m) => + m.status.value === ServerModelStatus.LOADED || + m.status.value === ServerModelStatus.SLEEPING + ) + .map((m) => m.id); + } + + get props() { + return this._props; + } + + get selectedModel(): ModelOption | null { + if (!this.selectedModelId) return null; + + return this.models.find((m) => m.id === this.selectedModelId) ?? null; + } + + get selectedModelContextSize(): number | null { + if (!this.selectedModelName) return null; + + return this.props.getModelContextSize(this.selectedModelName); + } + + /** + * Get model name in MODEL mode (single model). + * Extracts from model_path or model_alias from server props. + * In ROUTER mode, returns null (model is per-conversation). + */ + get singleModelName(): string | null { + if (serverStore.isRouterMode) return null; + + const props = serverStore.props; + + if (props?.model_alias) return props.model_alias; + + if (!props?.model_path) return null; + + return props.model_path.split(/(\\|\/)/).pop() || null; + } + + get status() { + return this._status; + } + + clearSelection(): void { + this.selectedModelId = null; + this.selectedModelName = null; + } + + /** + * Auto-selects the first available model if none is selected. + * Prioritizes: + * 1. Model from active conversation's last assistant response (if loaded) + * 2. Model from active conversation's last assistant response (if not loaded) + * 3. First loaded model (not from active conversation) + * 4. A favorite model + * 5. First available model + */ + async ensureFirstModelSelected(): Promise { + if (this.selectedModelName) return; + + const availableModels = this.getVisibleModels(); + + if (availableModels.length === 0) return; + + // Try to select model from last assistant response first + const lastModel = this.getModelFromLastAssistantResponse(); + + if (lastModel) { + const lastModelOption = availableModels.find((m) => m.model === lastModel); + + if (lastModelOption) { + await this.selectModelById(lastModelOption.id); + + if (this.isModelLoaded(lastModel)) { + await this.props.fetchModelProps(lastModel); + } + + return; + } + } + + // Try a loaded model first + const loadedModel = availableModels.find((m) => this.isModelLoaded(m.model)); + + if (loadedModel) { + await this.selectModelById(loadedModel.id); + await this.props.fetchModelProps(loadedModel.model); + + return; + } + + // Try loading a favorite model + const favorite = this.favoriteModelIds.values().next()?.value; + + if (favorite) { + await this.selectModelById(favorite); + + return; + } + + // Fall back to the first available model + await this.selectModelById(availableModels[0].id); + } + + /** + * Fetch list of models from server and detect server role. + * Also fetches modalities for MODEL mode (single model). + */ + async fetch(force = false): Promise { + if (this.inflightFetch) return this.inflightFetch; + + if (this.models.length > 0 && !force) return; + + this.inflightFetch = this.runFetch(); + try { + await this.inflightFetch; + } finally { + this.inflightFetch = null; + } + } + + /** + * Fetch router models with full metadata (ROUTER mode only). + * No-op in router mode — fetch() already calls listRouter() internally. + * Kept for API compatibility (e.g. handleOpenChange dropdown open handler). + */ + async fetchRouterModels(): Promise { + if (!serverStore.isRouterMode) return; + + try { + const response = await ModelsService.listRouter(); + + this.routerModels = response.data; + await this.props.fetchModalitiesForLoadedModels(); + + const visible = this.getVisibleModels(); + + if (visible.length === 1 && this.isModelLoaded(visible[0].model)) { + this.selectModelById(visible[0].id); + } + } catch (error) { + console.warn('Failed to fetch router models:', error); + this.routerModels = []; + } + } + + findModelById(modelId: string): ModelOption | null { + return this.models.find((model) => model.id === modelId) ?? null; + } + + findModelByName(modelName: string): ModelOption | null { + return ( + this.models.find( + (model) => + model.model === modelName || model.id === modelName || model.aliases?.includes(modelName) + ) ?? null + ); + } + + /** + * Gets the model name from the last assistant message in the active conversation. + * Used by both the chat page and settings page to maintain model consistency. + */ + getModelFromLastAssistantResponse(): string | null { + const messages = conversationsStore.activeMessages; + + if (!messages || messages.length === 0) return null; + + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].model) { + return messages[i].model; + } + } + + return null; + } + + getModelStatus(modelId: string): ServerModelStatus | null { + const model = this.routerModels.find((m) => m.id === modelId); + + return model?.status.value ?? null; + } + + hasModel(modelName: string): boolean { + return this.models.some((model) => model.model === modelName); + } + + isFavorite(modelId: string): boolean { + return this.favoriteModelIds.has(modelId); + } + + isModelLoaded(modelId: string): boolean { + const model = this.routerModels.find((m) => m.id === modelId); + + return ( + model?.status.value === ServerModelStatus.LOADED || + model?.status.value === ServerModelStatus.SLEEPING + ); + } + + async selectModelById(modelId: string): Promise { + if (!modelId || this.updating) return; + + if (this.selectedModelId === modelId) return; + + const option = this.models.find((model) => model.id === modelId); + + if (!option) throw new Error('Selected model is not available'); + + this.updating = true; + this.error = null; + + try { + this.selectedModelId = option.id; + this.selectedModelName = option.model; + } finally { + this.updating = false; + } + } + + /** + * Select a model by its model name (used for syncing with conversation model). + */ + selectModelByName(modelName: string): void { + const option = this.models.find((model) => model.model === modelName); + + if (option) { + this.selectedModelId = option.id; + this.selectedModelName = option.model; + } + } + + /** + * Auto-selects the model from the last assistant response if available and loaded. + * Returns true if a model was selected, false otherwise. + */ + async selectModelFromLastAssistantResponse(): Promise { + const lastModel = this.getModelFromLastAssistantResponse(); + + if (!lastModel || this.selectedModelName === lastModel) return false; + + const matchingModel = this.models.find((option) => option.model === lastModel); + + if (!matchingModel || !this.isModelLoaded(lastModel)) return false; + + try { + await this.selectModelById(matchingModel.id); + console.log(`[modelsStore] Automatically selected model: ${lastModel} from last message`); + + return true; + } catch (error) { + console.warn('[modelsStore] Failed to automatically select model from last message:', error); + + return false; + } + } + + toDisplayName(id: string): string { + const segments = id.split(/\\|\//); + const candidate = segments.pop(); + + return candidate && candidate.trim().length > 0 ? candidate : id; + } + + toggleFavorite(modelId: string): void { + const next = new SvelteSet(this.favoriteModelIds); + + if (next.has(modelId)) { + next.delete(modelId); + } else { + next.add(modelId); + } + + this.favoriteModelIds = next; + + try { + localStorage.setItem(FAVORITE_MODELS_LOCALSTORAGE_KEY, JSON.stringify([...next])); + } catch { + toast.error('Failed to save favorite models to local storage'); + } + } + + /** + * Build ModelOption[] from an API response. + * Both MODEL and ROUTER modes share the same mapping logic; + * they differ only in which endpoint is called. + */ + private buildModelOptions( + response: ApiModelListResponse | ApiRouterModelsListResponse + ): ModelOption[] { + return response.data.map((item: ApiModelDataEntry, index: number) => { + const details = response.models?.[index]; + const rawCapabilities = Array.isArray(details?.capabilities) ? details?.capabilities : []; + const displayNameSource = + details?.name && details.name.trim().length > 0 ? details.name : item.id; + const modelId = details?.model || item.id; + + return { + aliases: item.aliases ?? [], + capabilities: rawCapabilities.filter((value: unknown): value is string => Boolean(value)), + description: details?.description, + details: details?.details, + id: item.id, + meta: item.meta ?? null, + modalities: this.props.buildArchitectureModalities(item.architecture), + model: modelId, + name: this.toDisplayName(displayNameSource), + parsedId: ModelsService.parseModelId(modelId), + tags: item.tags ?? [] + }; + }); + } + + /** Fetch models in MODEL mode (single model, standard OpenAI-compatible). */ + private async fetchModelModeInternal(): Promise { + const response = await ModelsService.list(); + + return this.buildModelOptions(response); + } + + /** + * Filter to models visible in the UI (ui !== false). + */ + private getVisibleModels(): ModelOption[] { + return this.models.filter((option) => this.props.getModelProps(option.model)?.ui !== false); + } + + private loadFavoritesFromStorage(): Set { + try { + const raw = localStorage.getItem(FAVORITE_MODELS_LOCALSTORAGE_KEY); + + return raw ? new Set(JSON.parse(raw) as string[]) : new Set(); + } catch { + toast.error('Failed to load favorite models from local storage'); + + return new Set(); + } + } + + private async runFetch(): Promise { + this.loading = true; + this.error = null; + + try { + if (!serverStore.props) { + await serverStore.fetch(); + } + + const router = serverStore.isRouterMode; + + if (router) { + const response = await ModelsService.listRouter(); + + this.routerModels = response.data; + this.models = this.buildModelOptions(response); + + await this.props.fetchModalitiesForLoadedModels(); + + const visible = this.getVisibleModels(); + + if (visible.length === 1 && this.isModelLoaded(visible[0].model)) { + this.selectModelById(visible[0].id); + } + } else { + this.models = await this.fetchModelModeInternal(); + } + } catch (error) { + this.models = []; + this.error = error instanceof Error ? error.message : 'Failed to load models'; + + throw error; + } finally { + this.loading = false; + } + } +} + +export const modelsStore = new ModelsStore(); diff --git a/tools/ui/src/lib/stores/models/props.svelte.ts b/tools/ui/src/lib/stores/models/props.svelte.ts new file mode 100644 index 000000000000..9d2d817acb1d --- /dev/null +++ b/tools/ui/src/lib/stores/models/props.svelte.ts @@ -0,0 +1,273 @@ +/** + * ModelPropsManager - Per-model props cache, modalities and thinking detection + * + * Owns the /props?model= cache with TTL, the modality views over it, + * and chat-template thinking detection. Created and owned by modelsStore; + * the host owns the model lists that fetched modalities are mirrored onto. + * + * **API Inconsistency Workaround:** + * In MODEL mode, `/props` returns modalities for the single model. + * In ROUTER mode, `/props` has no modalities - must use `/props?model=` per model. + */ + +import { MODEL_PROPS_CACHE } from '$lib/constants'; +import { FileTypeCategory, ModelModality } from '$lib/enums'; +import { PropsService } from '$lib/services/props.service'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { serverStore } from '$lib/stores/server.svelte'; +// deep imports, not the '$lib/utils' barrel: it re-exports modules that reach back +// into the stores, and going through it here would read a half-built module +import { TTLCache } from '$lib/utils/cache-ttl'; +import { detectThinkingSupport } from '$lib/utils/chat-template-thinking-detector'; +import { SvelteSet } from 'svelte/reactivity'; + +/** + * The slice of modelsStore the manager reads. Kept narrow on purpose so it + * cannot reach around the host's full surface; modelsStore implements this + * structurally. + */ +export interface ModelPropsHost { + /** Model rows the manager mirrors fetched modalities onto. */ + models: ModelOption[]; + readonly selectedModelName: string | null; + readonly loadedModelIds: string[]; + isModelLoaded(modelId: string): boolean; +} + +export class ModelPropsManager { + /** Version counter for the cache - bumped on writes so $derived consumers recompute. */ + cacheVersion = $state(0); + /** + * Model-specific props cache with TTL. + * Key: modelId, Value: props data including modalities. + */ + private cache = new TTLCache({ + maxEntries: MODEL_PROPS_CACHE.MAX_ENTRIES, + ttlMs: MODEL_PROPS_CACHE.TTL_MS + }); + private fetching = new SvelteSet(); + + /** + * Whether the selected model's chat template supports thinking/reasoning. + * Uses heuristic detection on the model's chat_template from /props. + * + * - MODEL mode: the global /props already describes the single loaded model, + * so its chat_template is used directly and no per-model cache is involved + * - ROUTER mode: fetches /props?model= for the selected model (cached), + * triggering an async fetch if not yet cached + */ + get supportsThinking(): boolean { + if (!serverStore.isRouterMode) { + return detectThinkingSupport(serverStore.props?.chat_template ?? ''); + } + + const modelId = this.host.selectedModelName; + + if (!modelId) return false; + + if (!this.cache.get(modelId)) { + this.fetchModelProps(modelId); + } + + const props = this.getModelProps(modelId); + + return detectThinkingSupport(props?.chat_template ?? ''); + } + + /** Map the router modalities, the only source available while a model is not loaded. */ + buildArchitectureModalities( + architecture: ApiModelDataEntry['architecture'] + ): ModelModalities | undefined { + if (!architecture) return undefined; + + const inputs = architecture.input_modalities; + + return { + audio: inputs.includes(FileTypeCategory.AUDIO), + video: inputs.includes(FileTypeCategory.VIDEO), + vision: inputs.includes(FileTypeCategory.IMAGE) + }; + } + + /** + * Check if a specific model supports thinking. + * In MODEL mode the global /props describes the single loaded model. + * In ROUTER mode, fetches model props if not cached. + */ + checkModelSupportsThinking(modelId: string): boolean { + if (!serverStore.isRouterMode) { + return detectThinkingSupport(serverStore.props?.chat_template ?? ''); + } + + if (!modelId) return false; + + if (!this.cache.get(modelId)) { + this.fetchModelProps(modelId); + } + + const props = this.getModelProps(modelId); + + return detectThinkingSupport(props?.chat_template ?? ''); + } + + constructor(private host: ModelPropsHost) {} + + /** Fetch modalities for all loaded models from /props endpoint. */ + async fetchModalitiesForLoadedModels(): Promise { + const loadedModelIds = this.host.loadedModelIds; + + if (loadedModelIds.length === 0) return; + + const propsPromises = loadedModelIds.map((modelId) => this.fetchModelProps(modelId)); + + try { + const results = await Promise.all(propsPromises); + + this.host.models = this.host.models.map((model) => { + const modelIndex = loadedModelIds.indexOf(model.model); + + if (modelIndex === -1) return model; + + const props = results[modelIndex]; + + if (!props?.modalities) return model; + + return { ...model, modalities: this.buildModalities(props.modalities) }; + }); + + this.cacheVersion++; + } catch (error) { + console.warn('Failed to fetch modalities for loaded models:', error); + } + } + + /** + * Fetch props for a specific model from /props endpoint. + * Uses caching to avoid redundant requests. + * + * In ROUTER mode, this only fetches props if the model is loaded, + * since unloaded models return 400 from /props endpoint. + * + * @param modelId - Model identifier to fetch props for + * @returns Props data or null if fetch failed or model not loaded + */ + async fetchModelProps(modelId: string): Promise { + const cached = this.cache.get(modelId); + + if (cached) return cached; + + if (serverStore.isRouterMode && !this.host.isModelLoaded(modelId)) { + return null; + } + + if (this.fetching.has(modelId)) return null; + + this.fetching.add(modelId); + + try { + const props = await PropsService.fetchForModel(modelId); + + this.cache.set(modelId, props); + this.cacheVersion++; + + return props; + } catch (error) { + console.warn(`Failed to fetch props for model ${modelId}:`, error); + + return null; + } finally { + this.fetching.delete(modelId); + } + } + + getModelContextSize(modelId: string): number | null { + const props = this.getModelProps(modelId); + const nCtx = props?.default_generation_settings?.n_ctx; + + return typeof nCtx === 'number' ? nCtx : null; + } + + getModelModalities(modelId: string): ModelModalities | null { + if (!serverStore.isRouterMode && serverStore.props?.modalities) { + return this.buildModalities(serverStore.props.modalities); + } + + const model = this.host.models.find((m) => m.model === modelId || m.id === modelId); + + if (model?.modalities) { + return model.modalities; + } + + const props = this.cache.get(modelId); + + if (props?.modalities) { + return this.buildModalities(props.modalities); + } + + return null; + } + + getModelModalitiesArray(modelId: string): ModelModality[] { + const modalities = this.getModelModalities(modelId); + + if (!modalities) return []; + + const result: ModelModality[] = []; + + if (modalities.vision) result.push(ModelModality.VISION); + + if (modalities.audio) result.push(ModelModality.AUDIO); + + if (modalities.video) result.push(ModelModality.VIDEO); + + return result; + } + + getModelProps(modelId: string): ApiLlamaCppServerProps | null { + return this.cache.get(modelId); + } + + isModelPropsFetching(modelId: string): boolean { + return this.fetching.has(modelId); + } + + modelSupportsAudio(modelId: string): boolean { + return this.getModelModalities(modelId)?.audio ?? false; + } + + modelSupportsVideo(modelId: string): boolean { + return this.getModelModalities(modelId)?.video ?? false; + } + + modelSupportsVision(modelId: string): boolean { + return this.getModelModalities(modelId)?.vision ?? false; + } + + /** + * Update modalities for a specific model. + * Called when a model is loaded or when we need fresh modality data. + */ + async updateModelModalities(modelId: string): Promise { + const props = await this.fetchModelProps(modelId); + + if (!props?.modalities) return; + + this.host.models = this.host.models.map((model) => + model.model === modelId + ? { ...model, modalities: this.buildModalities(props.modalities!) } + : model + ); + + this.cacheVersion++; + } + + private buildModalities( + modalities: NonNullable + ): ModelModalities { + return { + audio: modalities.audio ?? false, + video: modalities.video ?? false, + vision: modalities.vision ?? false + }; + } +} diff --git a/tools/ui/src/lib/stores/models/status.svelte.ts b/tools/ui/src/lib/stores/models/status.svelte.ts new file mode 100644 index 000000000000..d0160aa4da3a --- /dev/null +++ b/tools/ui/src/lib/stores/models/status.svelte.ts @@ -0,0 +1,278 @@ +/** + * ModelStatusManager - Model load/unload operations and the /models/sse feed + * + * Owns the status feed subscription, load progress tracking, and the + * awaiters that settle load/unload operations. The feed drives status and + * progress, so it replaces any post-operation polling. Created and owned by + * modelsStore; the host owns the router model rows the feed updates. + */ + +import { ServerModelsSseEventType, ServerModelStatus } from '$lib/enums'; +import { ModelsService } from '$lib/services/models.service'; +import type { ModelPropsManager } from '$lib/stores/models/props.svelte'; +// direct imports between stores, not via the barrel, to avoid circular deps +import { serverStore } from '$lib/stores/server.svelte'; +import { SvelteMap } from 'svelte/reactivity'; +import { toast } from 'svelte-sonner'; + +/** + * The slice of modelsStore the manager drives. Kept narrow on purpose so it + * cannot reach around the host's full surface; modelsStore implements this + * structurally. + */ +export interface ModelStatusHost { + error: string | null; + readonly props: ModelPropsManager; + /** Router model rows the status feed updates. */ + routerModels: ApiModelDataEntry[]; + fetchRouterModels(): Promise; + isModelLoaded(modelId: string): boolean; + toDisplayName(id: string): string; +} + +export class ModelStatusManager { + private loadingStates = new SvelteMap(); + private loadProgress = new SvelteMap(); + // /models/sse feed state, the single source of truth for status and load progress + private statusAbort: AbortController | null = null; + private statusReaderActive = false; + private statusWaiters = new SvelteMap< + string, + { target: ServerModelStatus; resolve: () => void; reject: (e: Error) => void } + >(); + + constructor(private host: ModelStatusHost) {} + + async ensureLoaded(modelId: string): Promise { + if (this.host.isModelLoaded(modelId)) return; + + await this.load(modelId); + } + + /** + * Current load progress for a model, or null when not loading. + */ + getLoadProgress(modelId: string): ModelLoadProgress | null { + return this.loadProgress.get(modelId) ?? null; + } + + isOperationInProgress(modelId: string): boolean { + return this.loadingStates.get(modelId) ?? false; + } + + async load(modelId: string): Promise { + if (this.host.isModelLoaded(modelId)) return; + + if (this.loadingStates.get(modelId)) return; + + this.loadingStates.set(modelId, true); + this.host.error = null; + + // the feed drives completion, so it must be live before the request + this.subscribe(); + + const reachedLoaded = this.waitForStatus(modelId, ServerModelStatus.LOADED); + + reachedLoaded.catch(() => {}); + + try { + await ModelsService.load(modelId); + await reachedLoaded; + toast.success(`Model loaded: ${this.host.toDisplayName(modelId)}`); + } catch (error) { + this.rejectStatus(modelId, error instanceof Error ? error : new Error('load failed')); + this.host.error = error instanceof Error ? error.message : 'Failed to load model'; + toast.error(`Failed to load model: ${this.host.toDisplayName(modelId)}`); + + throw error; + } finally { + this.loadingStates.set(modelId, false); + } + } + + /** + * Open the /models/sse feed and keep it live with auto reconnect. + * Idempotent and router mode only. + */ + subscribe(): void { + if (this.statusReaderActive) return; + + if (!serverStore.isRouterMode) return; + + this.statusReaderActive = true; + this.statusAbort = new AbortController(); + void this.runStatusReader(this.statusAbort.signal); + } + + async unload(modelId: string): Promise { + if (!this.host.isModelLoaded(modelId)) return; + + if (this.loadingStates.get(modelId)) return; + + this.loadingStates.set(modelId, true); + this.host.error = null; + + this.subscribe(); + + const reachedUnloaded = this.waitForStatus(modelId, ServerModelStatus.UNLOADED); + + reachedUnloaded.catch(() => {}); + + try { + await ModelsService.unload(modelId); + await reachedUnloaded; + toast.info(`Model unloaded: ${this.host.toDisplayName(modelId)}`); + } catch (error) { + this.rejectStatus(modelId, error instanceof Error ? error : new Error('unload failed')); + this.host.error = error instanceof Error ? error.message : 'Failed to unload model'; + toast.error(`Failed to unload model: ${this.host.toDisplayName(modelId)}`); + + throw error; + } finally { + this.loadingStates.set(modelId, false); + } + } + + /** + * Close the /models/sse feed and drop transient progress. + */ + unsubscribe(): void { + this.statusReaderActive = false; + this.statusAbort?.abort(); + this.statusAbort = null; + this.loadProgress.clear(); + } + + /** + * Apply a status envelope: update the model row, track or clear progress, + * settle any pending load or unload awaiter. + */ + private applyModelStatus(event: ApiModelsSseEvent): void { + const model = event.model; + const data = event.data; + + if (!model || !data?.status) return; + + const status = data.status; + + this.setRouterModelStatus(model, status); + + if (status === ServerModelStatus.LOADING) { + if (data.progress) this.loadProgress.set(model, data.progress); + } else { + this.loadProgress.delete(model); + } + + if (status === ServerModelStatus.LOADED) { + void this.host.props.updateModelModalities(model); + } + + const failed = + status === ServerModelStatus.FAILED || + (status === ServerModelStatus.UNLOADED && (data.exit_code ?? 0) !== 0); + + if (failed) { + this.rejectStatus(model, new Error(`Model failed: ${this.host.toDisplayName(model)}`)); + + return; + } + + this.settleStatus(model, status); + } + + /** + * Route one feed record by event kind. Only the status_* events carry a + * status payload, models_reload triggers a list refresh, model_remove drops + * the row, download_* belong to the download surface, not here. + */ + private applyStatusEvent(event: ApiModelsSseEvent): void { + switch (event.event) { + case ServerModelsSseEventType.STATUS_CHANGE: + case ServerModelsSseEventType.MODEL_STATUS: + case ServerModelsSseEventType.STATUS_UPDATE: + this.applyModelStatus(event); + + break; + case ServerModelsSseEventType.MODELS_RELOAD: + void this.host.fetchRouterModels(); + + break; + case ServerModelsSseEventType.MODEL_REMOVE: + this.removeRouterModel(event.model); + + break; + case ServerModelsSseEventType.DOWNLOAD_PROGRESS: + break; + } + } + + /** + * Reject and drop the awaiter for a model. + */ + private rejectStatus(modelId: string, error: Error): void { + const waiter = this.statusWaiters.get(modelId); + + if (waiter) { + this.statusWaiters.delete(modelId); + waiter.reject(error); + } + } + + /** + * Drop a model row reported gone by the feed and settle its awaiters. + */ + private removeRouterModel(modelId: string): void { + if (this.host.routerModels.findIndex((m) => m.id === modelId) === -1) return; + + this.host.routerModels = this.host.routerModels.filter((m) => m.id !== modelId); + this.loadProgress.delete(modelId); + this.rejectStatus(modelId, new Error(`Model removed: ${this.host.toDisplayName(modelId)}`)); + } + + /** + * Read the feed and reconnect until unsubscribed. + */ + private async runStatusReader(signal: AbortSignal): Promise { + await ModelsService.watchModelEvents(signal, (event) => this.applyStatusEvent(event)); + } + + /** + * Update one model row status in place, reassigning to trigger reactivity. + */ + private setRouterModelStatus(modelId: string, status: ServerModelStatus): void { + const idx = this.host.routerModels.findIndex((m) => m.id === modelId); + + if (idx === -1) return; + + const current = this.host.routerModels[idx]; + + if (current.status.value === status) return; + + const next = [...this.host.routerModels]; + + next[idx] = { ...current, status: { ...current.status, value: status } }; + this.host.routerModels = next; + } + + /** + * Resolve and drop the awaiter when the model reaches its target status. + */ + private settleStatus(modelId: string, status: ServerModelStatus): void { + const waiter = this.statusWaiters.get(modelId); + + if (waiter && waiter.target === status) { + this.statusWaiters.delete(modelId); + waiter.resolve(); + } + } + + /** + * Register an awaiter that resolves when the feed reports target status. + * One operation runs per model at a time, so one awaiter per model is kept. + */ + private waitForStatus(modelId: string, target: ServerModelStatus): Promise { + return new Promise((resolve, reject) => { + this.statusWaiters.set(modelId, { reject, resolve, target }); + }); + } +} diff --git a/tools/ui/src/lib/stores/permissions.svelte.ts b/tools/ui/src/lib/stores/permissions.svelte.ts index 3e83538e9543..f4eae4b7e6a4 100644 --- a/tools/ui/src/lib/stores/permissions.svelte.ts +++ b/tools/ui/src/lib/stores/permissions.svelte.ts @@ -1,3 +1,11 @@ +/** + * permissionsStore - Allowed tool permissions + * + * Owns the set of tools the user has permanently allowed, persisted to + * localStorage. The agentic loop's permission gates consult it to run a + * tool without prompting. + */ + import { browser } from '$app/environment'; import { ALWAYS_ALLOWED_TOOLS_LOCALSTORAGE_KEY } from '$lib/constants'; import { SvelteSet } from 'svelte/reactivity'; @@ -5,6 +13,24 @@ import { SvelteSet } from 'svelte/reactivity'; class PermissionsStore { private _tools = $state(new SvelteSet()); + get tools(): ReadonlySet { + return this._tools; + } + + allowTool(key: string): void { + this._tools.add(key); + this.persist(); + } + + allowTools(keys: string[]): void { + for (const key of keys) this._tools.add(key); + this.persist(); + } + + hasTool(key: string): boolean { + return this._tools.has(key); + } + /** * Load persisted permissions. Called by initStores() after migrations * have run. @@ -29,30 +55,12 @@ class PermissionsStore { } } - get tools(): ReadonlySet { - return this._tools; - } - - hasTool(key: string): boolean { - return this._tools.has(key); - } - - allowTool(key: string): void { - this._tools.add(key); - this._persist(); - } - - allowTools(keys: string[]): void { - for (const key of keys) this._tools.add(key); - this._persist(); - } - revokeTool(key: string): void { this._tools.delete(key); - this._persist(); + this.persist(); } - private _persist(): void { + private persist(): void { try { localStorage.setItem(ALWAYS_ALLOWED_TOOLS_LOCALSTORAGE_KEY, JSON.stringify([...this._tools])); } catch (err) { diff --git a/tools/ui/src/lib/stores/server.svelte.ts b/tools/ui/src/lib/stores/server.svelte.ts index 7de5850b9eff..e145e2891dd7 100644 --- a/tools/ui/src/lib/stores/server.svelte.ts +++ b/tools/ui/src/lib/stores/server.svelte.ts @@ -1,79 +1,57 @@ +/** + * serverStore - Server connection state, configuration and role detection + * + * Owns the connection state and properties fetched from /props, plus MODEL + * vs ROUTER role detection and server-wide generation defaults. Uses + * PropsService for the /props fetch. + */ + import { ServerRole } from '$lib/enums'; import { PropsService } from '$lib/services/props.service'; import { ApiError } from '$lib/utils'; const LOADING_RETRY_INTERVAL_MS = 1000; -/** - * serverStore - Server connection state, configuration, and role detection - * - * This store manages the server connection state and properties fetched from `/props`. - * It provides reactive state for server configuration and role detection. - * - * **Architecture & Relationships:** - * - **PropsService**: Stateless service for fetching `/props` data - * - **serverStore** (this class): Reactive store for server state - * - **modelsStore**: Independent store for model management (uses PropsService directly) - * - * **Key Features:** - * - **Server State**: Connection status, loading, error handling - * - **Role Detection**: MODEL (single model) vs ROUTER (multi-model) - * - **Default Params**: Server-wide generation defaults - */ class ServerStore { - /** - * - * - * State - * - * - */ - - props = $state(null); - loading = $state(false); error = $state(null); - status = $state(null); + loading = $state(false); + props = $state(null); role = $state(null); + status = $state(null); private fetchPromise: Promise | null = null; private retryTimer: ReturnType | null = null; - /** - * - * - * Getters - * - * - */ - - get defaultParams(): ApiLlamaCppServerProps['default_generation_settings']['params'] | null { - return this.props?.default_generation_settings?.params || null; - } - get contextSize(): number | null { const nCtx = this.props?.default_generation_settings?.n_ctx; return typeof nCtx === 'number' ? nCtx : null; } - get uiSettings(): Record | undefined { - return this.props?.ui_settings ?? this.props?.webui_settings; + get defaultParams(): ApiLlamaCppServerProps['default_generation_settings']['params'] | null { + return this.props?.default_generation_settings?.params || null; + } + + get isModelMode(): boolean { + return this.role === ServerRole.MODEL; } get isRouterMode(): boolean { return this.role === ServerRole.ROUTER; } - get isModelMode(): boolean { - return this.role === ServerRole.MODEL; + get uiSettings(): Record | undefined { + return this.props?.ui_settings ?? this.props?.webui_settings; } - /** - * - * - * Data Handling - * - * - */ + clear(): void { + this.clearRetryTimer(); + this.props = null; + this.error = null; + this.status = null; + this.loading = false; + this.role = null; + this.fetchPromise = null; + } /** * @param background - Set by the automatic "still loading" poll. Skips the @@ -124,25 +102,6 @@ class ServerStore { await fetchPromise; } - clear(): void { - this.clearRetryTimer(); - this.props = null; - this.error = null; - this.status = null; - this.loading = false; - this.role = null; - this.fetchPromise = null; - } - - private scheduleRetry(): void { - if (this.retryTimer) return; - - this.retryTimer = setTimeout(() => { - this.retryTimer = null; - this.fetch({ background: true }); - }, LOADING_RETRY_INTERVAL_MS); - } - private clearRetryTimer(): void { if (this.retryTimer) { clearTimeout(this.retryTimer); @@ -150,14 +109,6 @@ class ServerStore { } } - /** - * - * - * Utilities - * - * - */ - private detectRole(props: ApiLlamaCppServerProps): void { const newRole = props?.role === ServerRole.ROUTER ? ServerRole.ROUTER : ServerRole.MODEL; @@ -166,6 +117,15 @@ class ServerStore { console.info(`Server running in ${newRole === ServerRole.ROUTER ? 'ROUTER' : 'MODEL'} mode`); } } + + private scheduleRetry(): void { + if (this.retryTimer) return; + + this.retryTimer = setTimeout(() => { + this.retryTimer = null; + this.fetch({ background: true }); + }, LOADING_RETRY_INTERVAL_MS); + } } export const serverStore = new ServerStore(); diff --git a/tools/ui/src/lib/stores/settings.svelte.ts b/tools/ui/src/lib/stores/settings/index.svelte.ts similarity index 89% rename from tools/ui/src/lib/stores/settings.svelte.ts rename to tools/ui/src/lib/stores/settings/index.svelte.ts index f23f6953ae27..0373ade42028 100644 --- a/tools/ui/src/lib/stores/settings.svelte.ts +++ b/tools/ui/src/lib/stores/settings/index.svelte.ts @@ -1,34 +1,10 @@ /** * settingsStore - Application configuration and theme management * - * This store manages all application settings including AI model parameters, UI preferences, - * and theme configuration. It provides persistent storage through localStorage with reactive - * state management using Svelte 5 runes. - * - * **Architecture & Relationships:** - * - **settingsStore** (this class): Configuration state management - * - Manages AI model parameters (temperature, max tokens, etc.) - * - Handles theme switching and persistence - * - Provides localStorage synchronization - * - Offers reactive configuration access - * - * - **ChatService**: Reads model parameters for API requests - * - **UI Components**: Subscribe to theme and configuration changes - * - * **Key Features:** - * - **Model Parameters**: Temperature, max tokens, top-p, top-k, repeat penalty - * - **Theme Management**: Auto, light, dark theme switching - * - **Persistence**: Automatic localStorage synchronization - * - **Reactive State**: Svelte 5 runes for automatic UI updates - * - **Default Handling**: Graceful fallback to defaults for missing settings - * - **Batch Updates**: Efficient multi-setting updates - * - **Reset Functionality**: Restore defaults for individual or all settings - * - * **Configuration Categories:** - * - Generation parameters (temperature, tokens, sampling) - * - UI preferences (theme, display options) - * - System settings (model selection, prompts) - * - Advanced options (seed, penalties, context handling) + * Owns generation parameters, UI preferences and theme, persisted to + * localStorage with Svelte 5 runes. Applies the admin's server ui_settings + * as defaults on first visit; sampling parameters sync with the server via + * ParameterSyncService. */ import { browser } from '$app/environment'; @@ -53,14 +29,6 @@ import { import { setMode } from 'mode-watcher'; class SettingsStore { - /** - * - * - * State - * - * - */ - config = $state({ ...SETTING_CONFIG_DEFAULT }); isInitialized = $state(false); userOverrides = $state>(new Set()); @@ -69,208 +37,210 @@ class SettingsStore { // application of server ui_settings defaults for new users. private isFirstVisit = false; + canSyncParameter(key: string): boolean { + return ParameterSyncService.canSyncParameter(key); + } /** - * - * - * Utilities (private helpers) - * - * - */ - - /** - * Helper method to get server defaults with null safety - * Centralizes the pattern of getting and extracting server defaults + * Clear all user overrides (for debugging) */ - private getServerDefaults(): Record { - return ParameterSyncService.extractServerDefaults(serverStore.defaultParams); + clearAllUserOverrides(): void { + this.userOverrides.clear(); + this.saveConfig(); + console.log('Cleared all user overrides'); } /** - * - * - * Lifecycle - * - * + * Export all settings as a versioned JSON-compatible object. + * The export captures the full config (excluding sensitive values like API key) + * and user overrides. Sensitive fields are filtered out for security by default. + * @param includeSensitiveData - If true, include sensitive fields (apiKey, MCP server headers) in export */ + exportSettings(includeSensitiveData: boolean = false): SettingsExportType { + // Build config excluding sensitive data unless user opts in + const configToExport: Record = + includeSensitiveData + ? { ...this.config } + : Object.fromEntries(Object.entries(this.config).filter(([key]) => key !== 'apiKey')); - /** - * Initialize the settings store by loading from localStorage. - * Called by initStores() after migrations have run. - */ - initialize() { - if (!browser) return; + // Handle MCP servers: exclude custom headers unless user opts in + if ('mcpServers' in configToExport && !includeSensitiveData) { + try { + const mcpServers = JSON.parse(configToExport.mcpServers as string) as Array< + Record + >; + const safeServers = mcpServers.map((server) => { + delete server.headers; - try { - this.loadConfig(); - this.migrateLegacyTheme(); - // Apply the persisted theme from config on initial load - setMode(this.config[SETTINGS_KEYS.THEME] as ColorMode); - this.isInitialized = true; - } catch (error) { - console.error('Failed to initialize settings store:', error); + return server; + }); + + configToExport.mcpServers = JSON.stringify(safeServers); + } catch { + // If parsing fails, just exclude the entire mcpServers field + delete (configToExport as Record).mcpServers; + } } + + return { + config: configToExport, + timestamp: Date.now(), + userOverrides: Array.from(this.userOverrides), + version: 1 + }; } /** - * Load configuration from localStorage - * Returns default values for missing keys to prevent breaking changes + * Reset all parameters to their default values (from props) + * This is used by the "Reset to Default" functionality + * Prioritizes Server defaults from /props, falls back to UI defaults */ - private loadConfig() { - if (!browser) return; + forceSyncWithServerDefaults(): void { + const propsDefaults = this.getServerDefaults(); + const uiSettings = serverStore.uiSettings; - try { - const storedConfigRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY); + for (const key of ParameterSyncService.getSyncableParameterKeys()) { + if (uiSettings && key in uiSettings) { + // UI setting from admin config: write actual value + setConfigValue(this.config, key, uiSettings[key]); + } else if (propsDefaults[key] !== undefined) { + // sampling param: clear it, let server decide + setConfigValue(this.config, key, ''); + } else if (key in SETTING_CONFIG_DEFAULT) { + setConfigValue(this.config, key, getConfigValue(SETTING_CONFIG_DEFAULT, key)); + } - // First visit: no stored config yet. Server ui_settings apply once in - // this state, then the user's config diverges freely. - this.isFirstVisit = storedConfigRaw === null; + this.userOverrides.delete(key); + } - const savedVal = JSON.parse(storedConfigRaw || '{}'); + // Non-syncable keys: reset is a full return to the instance state, the + // admin baseline value when defined, the factory default otherwise. + for (const key of Object.keys(SETTING_CONFIG_DEFAULT)) { + if (ParameterSyncService.canSyncParameter(key)) { + continue; + } - // Merge with defaults to prevent breaking changes - this.config = { - ...SETTING_CONFIG_DEFAULT, - ...savedVal - }; + const value = + uiSettings && key in uiSettings && uiSettings[key] !== undefined + ? uiSettings[key] + : getConfigValue(SETTING_CONFIG_DEFAULT, key); - // Default sendOnEnter to false on mobile when the user has no saved preference - if (!(SETTINGS_KEYS.SEND_ON_ENTER in savedVal)) { - if (deviceStore.isMobile) { - this.config[SETTINGS_KEYS.SEND_ON_ENTER] = false; - } - } + setConfigValue(this.config, key, value); - // Load user overrides - const savedOverrides = JSON.parse( - localStorage.getItem(USER_OVERRIDES_LOCALSTORAGE_KEY) || '[]' - ); + if (key === SETTINGS_KEYS.THEME) { + setMode(value as ColorMode); + } - this.userOverrides = new Set(savedOverrides); - } catch (error) { - console.warn('Failed to parse config from localStorage, using defaults:', error); - this.config = { ...SETTING_CONFIG_DEFAULT }; - this.userOverrides = new Set(); + this.userOverrides.delete(key); } + + this.saveConfig(); } /** - * Migrate the legacy un-namespaced "theme" localStorage key into config. - * Previously theme was stored separately in localStorage("theme") — now it lives - * inside the config object alongside all other settings. - * After migration the legacy key is removed. + * Get the entire configuration object + * @returns The complete configuration object */ - private migrateLegacyTheme() { - if (!browser) return; - - const legacyTheme = localStorage.getItem('theme'); - - if (legacyTheme) { - this.config[SETTINGS_KEYS.THEME] = legacyTheme; - localStorage.removeItem('theme'); - this.saveConfig(); - setMode(legacyTheme as ColorMode); - } + getAllConfig(): SettingsConfigType { + return { ...this.config }; } + /** - * - * - * Config Updates - * - * + * Get a specific configuration value + * @param key - The configuration key to get + * @returns The configuration value */ + getConfig(key: K): SettingsConfigType[K] { + return this.config[key]; + } /** - * Update a specific configuration setting - * @param key - The configuration key to update - * @param value - The new value for the configuration key + * Get diff between current settings and server defaults */ - updateConfig(key: K, value: SettingsConfigType[K]): void { - this.config[key] = value; - - if (ParameterSyncService.canSyncParameter(key as string)) { - const propsDefaults = this.getServerDefaults(); - const propsDefault = propsDefaults[key as string]; + getParameterDiff() { + const serverDefaults = this.getServerDefaults(); - if (propsDefault !== undefined) { - const normalizedValue = normalizeFloatingPoint(value); - const normalizedDefault = normalizeFloatingPoint(propsDefault); + if (Object.keys(serverDefaults).length === 0) return {}; - if (normalizedValue === normalizedDefault) { - this.userOverrides.delete(key as string); - } else { - this.userOverrides.add(key as string); - } - } - } + const configAsRecord = configToParameterRecord( + this.config, + ParameterSyncService.getSyncableParameterKeys() + ); - this.saveConfig(); + return ParameterSyncService.createParameterDiff(configAsRecord, serverDefaults); } /** - * Update multiple configuration settings at once - * @param updates - Object containing the configuration updates + * Get parameter information including source for a specific parameter */ - updateMultipleConfig(updates: Partial) { - Object.assign(this.config, updates); - + getParameterInfo(key: string) { const propsDefaults = this.getServerDefaults(); + const currentValue = getConfigValue(this.config, key); - for (const [key, value] of Object.entries(updates)) { - if (ParameterSyncService.canSyncParameter(key)) { - const propsDefault = propsDefaults[key]; + return ParameterSyncService.getParameterInfo( + key, + currentValue ?? '', + propsDefaults, + this.userOverrides + ); + } - if (propsDefault !== undefined) { - const normalizedValue = normalizeFloatingPoint(value); - const normalizedDefault = normalizeFloatingPoint(propsDefault); + /** + * Import settings from a previously exported object. + * Restores config (including theme) and user overrides. + * @param data - The exported settings object + */ + importSettings(data: SettingsExportType): void { + if (!browser) return; - if (normalizedValue === normalizedDefault) { - this.userOverrides.delete(key); - } else { - this.userOverrides.add(key); - } - } - } + if (!data || !data.config) { + throw new Error('Invalid settings data: missing config'); } + // Restore config (theme is included in config) + this.config = { + ...SETTING_CONFIG_DEFAULT, + ...data.config + }; + + // Restore user overrides (derived state — may be stale if server defaults differ) + this.userOverrides = new Set(data.userOverrides ?? []); + + // Persist to localStorage this.saveConfig(); + + // Apply theme for immediate visual feedback + setMode(this.config[SETTINGS_KEYS.THEME] as ColorMode); + + console.log('Settings imported successfully'); } /** - * Save the current configuration to localStorage + * Initialize the settings store by loading from localStorage. + * Called by initStores() after migrations have run. */ - private saveConfig() { + initialize() { if (!browser) return; try { - localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(this.config)); - - localStorage.setItem( - USER_OVERRIDES_LOCALSTORAGE_KEY, - JSON.stringify(Array.from(this.userOverrides)) - ); + this.loadConfig(); + this.migrateLegacyTheme(); + // Apply the persisted theme from config on initial load + setMode(this.config[SETTINGS_KEYS.THEME] as ColorMode); + this.isInitialized = true; } catch (error) { - console.error('Failed to save config to localStorage:', error); + console.error('Failed to initialize settings store:', error); } } /** - * Update the theme setting. - * @param newTheme - The new theme value + * Reset all settings to defaults. */ - updateTheme(newTheme: string) { - this.updateConfig(SETTINGS_KEYS.THEME, newTheme); + resetAll() { + this.resetConfig(); - setMode(newTheme as ColorMode); + this.resetTheme(); } - /** - * - * - * Reset - * - * - */ - /** * Reset configuration to defaults */ @@ -280,25 +250,6 @@ class SettingsStore { this.saveConfig(); } - /** - * Reset theme to default value. - * Theme is now stored inside the config object. - */ - resetTheme() { - this.updateConfig(SETTINGS_KEYS.THEME, SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.THEME]); - - setMode(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.THEME] as ColorMode); - } - - /** - * Reset all settings to defaults. - */ - resetAll() { - this.resetConfig(); - - this.resetTheme(); - } - /** * Reset a parameter to Server default (or UI default if no Server default) */ @@ -316,18 +267,20 @@ class SettingsStore { setConfigValue(this.config, key, getConfigValue(SETTING_CONFIG_DEFAULT, key)); } - this.userOverrides.delete(key); - this.saveConfig(); + this.userOverrides.delete(key); + this.saveConfig(); + } + + /** + * Reset theme to default value. + * Theme is now stored inside the config object. + */ + resetTheme() { + this.updateConfig(SETTINGS_KEYS.THEME, SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.THEME]); + + setMode(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.THEME] as ColorMode); } - /** - * - * - * Server Sync - * - * - */ - /** * Initialize settings with props defaults when server properties are first loaded * This sets up the default values from /props endpoint @@ -385,47 +338,27 @@ class SettingsStore { } /** - * Reset all parameters to their default values (from props) - * This is used by the "Reset to Default" functionality - * Prioritizes Server defaults from /props, falls back to UI defaults + * Update a specific configuration setting + * @param key - The configuration key to update + * @param value - The new value for the configuration key */ - forceSyncWithServerDefaults(): void { - const propsDefaults = this.getServerDefaults(); - const uiSettings = serverStore.uiSettings; - - for (const key of ParameterSyncService.getSyncableParameterKeys()) { - if (uiSettings && key in uiSettings) { - // UI setting from admin config: write actual value - setConfigValue(this.config, key, uiSettings[key]); - } else if (propsDefaults[key] !== undefined) { - // sampling param: clear it, let server decide - setConfigValue(this.config, key, ''); - } else if (key in SETTING_CONFIG_DEFAULT) { - setConfigValue(this.config, key, getConfigValue(SETTING_CONFIG_DEFAULT, key)); - } - - this.userOverrides.delete(key); - } - - // Non-syncable keys: reset is a full return to the instance state, the - // admin baseline value when defined, the factory default otherwise. - for (const key of Object.keys(SETTING_CONFIG_DEFAULT)) { - if (ParameterSyncService.canSyncParameter(key)) { - continue; - } + updateConfig(key: K, value: SettingsConfigType[K]): void { + this.config[key] = value; - const value = - uiSettings && key in uiSettings && uiSettings[key] !== undefined - ? uiSettings[key] - : getConfigValue(SETTING_CONFIG_DEFAULT, key); + if (ParameterSyncService.canSyncParameter(key as string)) { + const propsDefaults = this.getServerDefaults(); + const propsDefault = propsDefaults[key as string]; - setConfigValue(this.config, key, value); + if (propsDefault !== undefined) { + const normalizedValue = normalizeFloatingPoint(value); + const normalizedDefault = normalizeFloatingPoint(propsDefault); - if (key === SETTINGS_KEYS.THEME) { - setMode(value as ColorMode); + if (normalizedValue === normalizedDefault) { + this.userOverrides.delete(key as string); + } else { + this.userOverrides.add(key as string); + } } - - this.userOverrides.delete(key); } this.saveConfig(); @@ -434,148 +367,143 @@ class SettingsStore { /** * * - * Utilities + * Import / Export * * */ /** - * Get a specific configuration value - * @param key - The configuration key to get - * @returns The configuration value - */ - getConfig(key: K): SettingsConfigType[K] { - return this.config[key]; - } - - /** - * Get the entire configuration object - * @returns The complete configuration object + * Update multiple configuration settings at once + * @param updates - Object containing the configuration updates */ - getAllConfig(): SettingsConfigType { - return { ...this.config }; - } - - canSyncParameter(key: string): boolean { - return ParameterSyncService.canSyncParameter(key); - } + updateMultipleConfig(updates: Partial) { + Object.assign(this.config, updates); - /** - * Get parameter information including source for a specific parameter - */ - getParameterInfo(key: string) { const propsDefaults = this.getServerDefaults(); - const currentValue = getConfigValue(this.config, key); - - return ParameterSyncService.getParameterInfo( - key, - currentValue ?? '', - propsDefaults, - this.userOverrides - ); - } - /** - * Get diff between current settings and server defaults - */ - getParameterDiff() { - const serverDefaults = this.getServerDefaults(); + for (const [key, value] of Object.entries(updates)) { + if (ParameterSyncService.canSyncParameter(key)) { + const propsDefault = propsDefaults[key]; - if (Object.keys(serverDefaults).length === 0) return {}; + if (propsDefault !== undefined) { + const normalizedValue = normalizeFloatingPoint(value); + const normalizedDefault = normalizeFloatingPoint(propsDefault); - const configAsRecord = configToParameterRecord( - this.config, - ParameterSyncService.getSyncableParameterKeys() - ); + if (normalizedValue === normalizedDefault) { + this.userOverrides.delete(key); + } else { + this.userOverrides.add(key); + } + } + } + } - return ParameterSyncService.createParameterDiff(configAsRecord, serverDefaults); + this.saveConfig(); } /** - * Clear all user overrides (for debugging) + * Update the theme setting. + * @param newTheme - The new theme value */ - clearAllUserOverrides(): void { - this.userOverrides.clear(); - this.saveConfig(); - console.log('Cleared all user overrides'); + updateTheme(newTheme: string) { + this.updateConfig(SETTINGS_KEYS.THEME, newTheme); + + setMode(newTheme as ColorMode); } /** * * - * Import / Export + * Utilities (private helpers) * * */ /** - * Export all settings as a versioned JSON-compatible object. - * The export captures the full config (excluding sensitive values like API key) - * and user overrides. Sensitive fields are filtered out for security by default. - * @param includeSensitiveData - If true, include sensitive fields (apiKey, MCP server headers) in export + * Helper method to get server defaults with null safety + * Centralizes the pattern of getting and extracting server defaults */ - exportSettings(includeSensitiveData: boolean = false): SettingsExportType { - // Build config excluding sensitive data unless user opts in - const configToExport: Record = - includeSensitiveData - ? { ...this.config } - : Object.fromEntries(Object.entries(this.config).filter(([key]) => key !== 'apiKey')); + private getServerDefaults(): Record { + return ParameterSyncService.extractServerDefaults(serverStore.defaultParams); + } - // Handle MCP servers: exclude custom headers unless user opts in - if ('mcpServers' in configToExport && !includeSensitiveData) { - try { - const mcpServers = JSON.parse(configToExport.mcpServers as string) as Array< - Record - >; - const safeServers = mcpServers.map((server) => { - delete server.headers; + /** + * Load configuration from localStorage + * Returns default values for missing keys to prevent breaking changes + */ + private loadConfig() { + if (!browser) return; - return server; - }); + try { + const storedConfigRaw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY); - configToExport.mcpServers = JSON.stringify(safeServers); - } catch { - // If parsing fails, just exclude the entire mcpServers field - delete (configToExport as Record).mcpServers; + // First visit: no stored config yet. Server ui_settings apply once in + // this state, then the user's config diverges freely. + this.isFirstVisit = storedConfigRaw === null; + + const savedVal = JSON.parse(storedConfigRaw || '{}'); + + // Merge with defaults to prevent breaking changes + this.config = { + ...SETTING_CONFIG_DEFAULT, + ...savedVal + }; + + // Default sendOnEnter to false on mobile when the user has no saved preference + if (!(SETTINGS_KEYS.SEND_ON_ENTER in savedVal)) { + if (deviceStore.isMobile) { + this.config[SETTINGS_KEYS.SEND_ON_ENTER] = false; + } } - } - return { - config: configToExport, - timestamp: Date.now(), - userOverrides: Array.from(this.userOverrides), - version: 1 - }; + // Load user overrides + const savedOverrides = JSON.parse( + localStorage.getItem(USER_OVERRIDES_LOCALSTORAGE_KEY) || '[]' + ); + + this.userOverrides = new Set(savedOverrides); + } catch (error) { + console.warn('Failed to parse config from localStorage, using defaults:', error); + this.config = { ...SETTING_CONFIG_DEFAULT }; + this.userOverrides = new Set(); + } } /** - * Import settings from a previously exported object. - * Restores config (including theme) and user overrides. - * @param data - The exported settings object + * Migrate the legacy un-namespaced "theme" localStorage key into config. + * Previously theme was stored separately in localStorage("theme") — now it lives + * inside the config object alongside all other settings. + * After migration the legacy key is removed. */ - importSettings(data: SettingsExportType): void { + private migrateLegacyTheme() { if (!browser) return; - if (!data || !data.config) { - throw new Error('Invalid settings data: missing config'); - } - - // Restore config (theme is included in config) - this.config = { - ...SETTING_CONFIG_DEFAULT, - ...data.config - }; + const legacyTheme = localStorage.getItem('theme'); - // Restore user overrides (derived state — may be stale if server defaults differ) - this.userOverrides = new Set(data.userOverrides ?? []); + if (legacyTheme) { + this.config[SETTINGS_KEYS.THEME] = legacyTheme; + localStorage.removeItem('theme'); + this.saveConfig(); + setMode(legacyTheme as ColorMode); + } + } - // Persist to localStorage - this.saveConfig(); + /** + * Save the current configuration to localStorage + */ + private saveConfig() { + if (!browser) return; - // Apply theme for immediate visual feedback - setMode(this.config[SETTINGS_KEYS.THEME] as ColorMode); + try { + localStorage.setItem(CONFIG_LOCALSTORAGE_KEY, JSON.stringify(this.config)); - console.log('Settings imported successfully'); + localStorage.setItem( + USER_OVERRIDES_LOCALSTORAGE_KEY, + JSON.stringify(Array.from(this.userOverrides)) + ); + } catch (error) { + console.error('Failed to save config to localStorage:', error); + } } } diff --git a/tools/ui/src/lib/stores/settings-referrer.svelte.ts b/tools/ui/src/lib/stores/settings/referrer.svelte.ts similarity index 50% rename from tools/ui/src/lib/stores/settings-referrer.svelte.ts rename to tools/ui/src/lib/stores/settings/referrer.svelte.ts index 297a0d6a4553..9679049df903 100644 --- a/tools/ui/src/lib/stores/settings-referrer.svelte.ts +++ b/tools/ui/src/lib/stores/settings/referrer.svelte.ts @@ -1,3 +1,10 @@ +/** + * settingsReferrer - Remembers the settings route to return to after exit + * + * Tracks the last settings section the user was on so the app can return + * there after a fallback exit. Standalone reactive value, no host. + */ + import { SETTINGS_FALLBACK_EXIT_ROUTE } from '$lib/constants'; let _url = $state(SETTINGS_FALLBACK_EXIT_ROUTE); diff --git a/tools/ui/src/lib/stores/tools.svelte.ts b/tools/ui/src/lib/stores/tools.svelte.ts index 9f044c83e67b..e255b8a43ec2 100644 --- a/tools/ui/src/lib/stores/tools.svelte.ts +++ b/tools/ui/src/lib/stores/tools.svelte.ts @@ -1,3 +1,12 @@ +/** + * toolsStore - Tool registry and enablement + * + * Owns the server tool listing (with working-directory resolution), built-in + * browser tools, MCP tools and per-tool enablement, exposed as a unified + * tool set for the LLM and the tools UI. Consumed by the agentic loop and + * the chat flows. + */ + import { browser } from '$app/environment'; import { buildBrowserInfoToolDefinition, @@ -18,9 +27,9 @@ import { } from '$lib/enums'; import { ToolsService } from '$lib/services/tools.service'; // direct imports between stores, not via the barrel, to avoid circular deps -import { mcpStore } from '$lib/stores/mcp.svelte'; -import { modelsStore } from '$lib/stores/models.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { mcpStore } from '$lib/stores/mcp/index.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { OpenAIToolDefinition, ToolEntry, ToolGroup } from '$lib/types'; import { buildSandboxToolDefinition } from '$lib/utils'; import { SvelteMap, SvelteSet } from 'svelte/reactivity'; @@ -28,156 +37,74 @@ import { SvelteMap, SvelteSet } from 'svelte/reactivity'; /** Stable selection identity for a tool, shared by the disabled set and the permission store */ class ToolsStore { - private _serverTools = $state([]); - private _loading = $state(false); - private _error = $state(null); private _disabledTools = $state(new SvelteSet()); + private _error = $state(null); + private _loading = $state(false); + private _serverHome = $state(undefined); + private _serverTools = $state([]); + private _toolsEndpointUnreachable = $state(false); // server tools that resolve their paths against the working directory, // as declared by the server in its `/tools` listing - private _cwdAwareTools = $state(new SvelteSet()); - private _toolsEndpointUnreachable = $state(false); - private _serverHome = $state(undefined); - - /** - * Load persisted disabled tools and fetch the builtin tool list. - * Called by initStores() after migrations have run. - */ - initialize(): void { - // browser-only init: skip on SSR to avoid localStorage/fetch side effects - if (!browser) return; - - try { - const stored = localStorage.getItem(DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY); - - if (stored) { - const parsed = JSON.parse(stored); - - if (Array.isArray(parsed)) { - for (const key of parsed) { - if (typeof key === 'string') this._disabledTools.add(key); - } - } - } - } catch (err) { - console.error('[ToolsStore] Failed to load disabled tools from localStorage:', err); - } - - this.fetchServerTools(); - } - - private persistDisabledTools(): void { - try { - localStorage.setItem( - DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY, - JSON.stringify([...this._disabledTools]) - ); - } catch { - // ignore storage errors - } - } - - private toolKey(source: ToolSource, name: string, serverId?: string): string { - switch (source) { - case ToolSource.MCP: - return serverId ? `mcp-${serverId}:${name}` : `mcp:${name}`; - case ToolSource.CUSTOM: - return `custom:${name}`; - case ToolSource.BROWSER: - return `browser:${name}`; - default: - return `server:${name}`; - } - } - - private inferTypeFromDefault(value: unknown): string | undefined { - if (typeof value === 'string') return 'string'; - - if (typeof value === 'boolean') return 'boolean'; - - if (typeof value === 'number') return Number.isInteger(value) ? 'integer' : 'number'; - - if (Array.isArray(value)) return 'array'; - - if (value !== null && typeof value === 'object') return 'object'; + private cwdAwareTools = $state(new SvelteSet()); - return undefined; + get allToolDefinitions(): OpenAIToolDefinition[] { + return this.allTools.map((t) => t.definition); } - /** - * Recursively normalize a JSON Schema object: infers `type` from `default` - * for properties / items that omit it, and descends into nested `properties` - * and `items`. Returns a new object -- does not mutate the input. - */ - private normalizeJsonSchema(schema: Record): Record { - if (!schema || typeof schema !== 'object') return schema; - - const normalized: Record = { ...schema }; - - if (normalized.properties && typeof normalized.properties === 'object') { - const props = normalized.properties as Record>; - const normalizedProps: Record> = {}; - - for (const [key, prop] of Object.entries(props)) { - if (!prop || typeof prop !== 'object') { - normalizedProps[key] = prop; - - continue; - } - - const normalizedProp: Record = { ...prop }; + /** Canonical flat list of tool entries with source metadata and stable keys, deduped by key */ + get allTools(): ToolEntry[] { + const entries: ToolEntry[] = []; + const seen = new SvelteSet(); + const push = (entry: ToolEntry) => { + if (seen.has(entry.key)) return; - if (!normalizedProp.type && normalizedProp.default !== undefined) { - const inferred = this.inferTypeFromDefault(normalizedProp.default); + seen.add(entry.key); + entries.push(entry); + }; - if (inferred) normalizedProp.type = inferred; - } + for (const def of this._serverTools) { + const name = def.function.name; - if (normalizedProp.properties) { - Object.assign( - normalizedProp, - this.normalizeJsonSchema(normalizedProp as Record) - ); - } + push({ + definition: def, + key: this.toolKey(ToolSource.SERVER, name), + source: ToolSource.SERVER + }); + } - if (normalizedProp.items && typeof normalizedProp.items === 'object') { - normalizedProp.items = this.normalizeJsonSchema( - normalizedProp.items as Record - ); - } + for (const def of this.browserTools) { + const name = def.function.name; - normalizedProps[key] = normalizedProp; - } - normalized.properties = normalizedProps; + push({ + definition: def, + key: this.toolKey(ToolSource.BROWSER, name), + source: ToolSource.BROWSER + }); } - return normalized; - } + for (const { definition, serverId, serverName } of this.mcpEntries()) { + const name = definition.function.name; - private mcpDefinition( - name: string, - description: string | undefined, - schema?: Record - ): OpenAIToolDefinition { - return { - function: { - description, - name, - parameters: schema ?? { properties: {}, required: [], type: JsonSchemaType.OBJECT } - }, - type: ToolCallType.FUNCTION - }; - } + push({ + definition, + key: this.toolKey(ToolSource.MCP, name, serverId), + serverId, + serverName, + source: ToolSource.MCP + }); + } - get serverTools(): OpenAIToolDefinition[] { - return this._serverTools; - } + for (const def of this.customTools) { + const name = def.function.name; - get serverHome(): string | null { - return this._serverHome ?? null; - } + push({ + definition: def, + key: this.toolKey(ToolSource.CUSTOM, name), + source: ToolSource.CUSTOM + }); + } - get mcpTools(): OpenAIToolDefinition[] { - return this.mcpEntries().map((e) => e.definition); + return entries; } get browserTools(): OpenAIToolDefinition[] { @@ -199,31 +126,6 @@ class ToolsStore { return tools; } - private hasServerTool(name: BuiltInTool): boolean { - return this._serverTools.some((def) => def.function.name === name); - } - - /** - * `read_media` runs in the browser on top of the server's `read_file`, so it - * exists only when that tool is served and the active model can perceive the - * bytes. The server cannot make this call - it does not know which model the - * conversation uses. - */ - private readMediaTool(): OpenAIToolDefinition | null { - if (!this.hasServerTool(BuiltInTool.SERVER_READ_FILE)) return null; - - const model = modelsStore.selectedModelName ?? modelsStore.models[0]?.model ?? ''; - - if (!model) return null; - - const vision = modelsStore.modelSupportsVision(model); - const audio = modelsStore.modelSupportsAudio(model); - - if (!vision && !audio) return null; - - return buildReadMediaToolDefinition(vision, audio); - } - get customTools(): OpenAIToolDefinition[] { const raw = settingsStore.config.customJson; @@ -248,109 +150,52 @@ class ToolsStore { } } - /** Normalize MCP tools from live connections when available, fall back to health check data */ - private mcpEntries(): { - serverId: string; - serverName: string; - definition: OpenAIToolDefinition; - }[] { - const out: { serverId: string; serverName: string; definition: OpenAIToolDefinition }[] = []; - const connections = mcpStore.getConnections(); + get disabledTools(): SvelteSet { + return this._disabledTools; + } - if (connections.size > 0) { - for (const [serverId, connection] of connections) { - const serverName = mcpStore.getServerDisplayName(serverId); + get error(): string | null { + return this._error; + } - for (const tool of connection.tools) { - const rawSchema = (tool.inputSchema as Record) ?? { - properties: {}, - required: [], - type: JsonSchemaType.OBJECT - }; + /** + * Check if a working directory is worth setting: at least one server tool + * that reads it is both served and left enabled by the user. + */ + get hasEnabledCwdTools(): boolean { + return this._serverTools.some((def) => { + const name = def.function.name; - out.push({ - definition: { - function: { - description: tool.description, - name: tool.name, - parameters: this.normalizeJsonSchema(rawSchema) - }, - type: ToolCallType.FUNCTION - }, - serverId, - serverName - }); - } - } - } else { - for (const { serverId, serverName, tools } of this.getMcpToolsFromHealthChecks()) { - for (const tool of tools) { - out.push({ - definition: this.mcpDefinition(tool.name, tool.description), - serverId, - serverName - }); - } - } - } - - return out; + return ( + this.cwdAwareTools.has(name) && + !this._disabledTools.has(this.toolKey(ToolSource.SERVER, name)) + ); + }); } - /** Canonical flat list of tool entries with source metadata and stable keys, deduped by key */ - get allTools(): ToolEntry[] { - const entries: ToolEntry[] = []; - const seen = new SvelteSet(); - const push = (entry: ToolEntry) => { - if (seen.has(entry.key)) return; - - seen.add(entry.key); - entries.push(entry); - }; - - for (const def of this._serverTools) { - const name = def.function.name; - - push({ - definition: def, - key: this.toolKey(ToolSource.SERVER, name), - source: ToolSource.SERVER - }); - } - - for (const def of this.browserTools) { - const name = def.function.name; - - push({ - definition: def, - key: this.toolKey(ToolSource.BROWSER, name), - source: ToolSource.BROWSER - }); - } + /** Check if there are any enabled tools available (server, MCP, or custom) */ + get hasEnabledTools(): boolean { + return this.getEnabledToolsForLLM().length > 0; + } - for (const { definition, serverId, serverName } of this.mcpEntries()) { - const name = definition.function.name; + get isToolsEndpointUnreachable(): boolean { + return this._toolsEndpointUnreachable; + } - push({ - definition, - key: this.toolKey(ToolSource.MCP, name, serverId), - serverId, - serverName, - source: ToolSource.MCP - }); - } + get loading(): boolean { + return this._loading; + } - for (const def of this.customTools) { - const name = def.function.name; + get mcpTools(): OpenAIToolDefinition[] { + return this.mcpEntries().map((e) => e.definition); + } - push({ - definition: def, - key: this.toolKey(ToolSource.CUSTOM, name), - source: ToolSource.CUSTOM - }); - } + get serverHome(): string | null { + return this._serverHome ?? null; + } - return entries; + get serverTools(): OpenAIToolDefinition[] { + return this._serverTools; } /** Tools grouped by category for tree display, derived from the canonical entries */ @@ -382,16 +227,47 @@ class ToolsStore { return groups; } - private groupLabel(entry: ToolEntry): string { - switch (entry.source) { - case ToolSource.MCP: - return entry.serverName ?? ''; - case ToolSource.CUSTOM: - return TOOL_GROUP_LABELS[ToolSource.CUSTOM]; - case ToolSource.BROWSER: - return TOOL_GROUP_LABELS[ToolSource.BROWSER]; - default: - return TOOL_GROUP_LABELS[ToolSource.SERVER]; + /** Enable all tools belonging to a specific MCP server */ + enableAllToolsForServer(serverId: string): void { + const connection = mcpStore.getConnections().get(serverId); + + if (!connection) return; + + for (const tool of connection.tools) { + this._disabledTools.delete(this.toolKey(ToolSource.MCP, tool.name, serverId)); + } + this.persistDisabledTools(); + } + + async fetchServerTools(): Promise { + if (this._loading) return; + + this._loading = true; + this._error = null; + this._toolsEndpointUnreachable = false; + + try { + const toolInfos = await ToolsService.list(); + + this._serverTools = toolInfos.map((info) => info.definition); + this.cwdAwareTools = new SvelteSet( + toolInfos.filter((info) => info.uses_cwd).map((info) => info.tool) + ); + } catch (err) { + const errorMessage = err instanceof Error ? err.message : String(err); + + this._error = errorMessage; + + // 403 from /tools means the server was started without --tools + // TODO: check status code instead of relying on message + if (errorMessage.includes('this feature is disabled')) { + this._toolsEndpointUnreachable = true; + console.info('[ToolsStore] Server tools are disabled on the server'); + } else { + console.error('[ToolsStore] Failed to fetch server tools:', err); + } + } finally { + this._loading = false; } } @@ -430,38 +306,92 @@ class ToolsStore { return result; } - get allToolDefinitions(): OpenAIToolDefinition[] { - return this.allTools.map((t) => t.definition); + /** Permission key for a tool name, identical to the selection key */ + getPermissionKey(toolName: string): string | null { + return this.findEntryByName(toolName)?.key ?? null; } - get loading(): boolean { - return this._loading; + /** Get the display label for the server that owns a given tool */ + getToolServerLabel(toolName: string): string { + const entry = this.findEntryByName(toolName); + + if (!entry) return ''; + + if (entry.serverName) return mcpStore.getServerDisplayName(entry.serverName); + + if (entry.source === ToolSource.SERVER) return TOOL_SERVER_LABELS[ToolSource.SERVER]; + + if (entry.source === ToolSource.CUSTOM) return TOOL_SERVER_LABELS[ToolSource.CUSTOM]; + + if (entry.source === ToolSource.BROWSER) return TOOL_SERVER_LABELS[ToolSource.BROWSER]; + + return ''; } - get error(): string | null { - return this._error; + /** Determine the source of a tool by its name */ + getToolSource(toolName: string): ToolSource | null { + return this.findEntryByName(toolName)?.source ?? null; } - get isToolsEndpointUnreachable(): boolean { - return this._toolsEndpointUnreachable; + /** + * Load persisted disabled tools and fetch the builtin tool list. + * Called by initStores() after migrations have run. + */ + initialize(): void { + // browser-only init: skip on SSR to avoid localStorage/fetch side effects + if (!browser) return; + + try { + const stored = localStorage.getItem(DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY); + + if (stored) { + const parsed = JSON.parse(stored); + + if (Array.isArray(parsed)) { + for (const key of parsed) { + if (typeof key === 'string') this._disabledTools.add(key); + } + } + } + } catch (err) { + console.error('[ToolsStore] Failed to load disabled tools from localStorage:', err); + } + + this.fetchServerTools(); } - get disabledTools(): SvelteSet { - return this._disabledTools; + isGroupFullyEnabled(group: ToolGroup): boolean { + return group.tools.length > 0 && group.tools.every((t) => this.isToolEnabled(t.key)); } isToolEnabled(key: string): boolean { return !this._disabledTools.has(key); } - toggleTool(key: string): void { - if (this._disabledTools.has(key)) { - this._disabledTools.delete(key); - } else { - this._disabledTools.add(key); + /** + * Absolute home directory on the server, resolved once per session via + * file_glob_search's `base` field (the server expands `~`). Anchors the + * directory picker's search scope and the `~` abbreviation of cwd + * displays. Returns null when tools are unavailable. + */ + async resolveServerHome(): Promise { + if (this._serverHome !== undefined) return this._serverHome; + + try { + const res = await ToolsService.executeToolRaw(BuiltInTool.SERVER_FILE_GLOB_SEARCH, { + limit: 1, + max_depth: 1, + path: HOME_TILDE, + type: GlobSearchType.DIR + }); + + this._serverHome = typeof res.base === 'string' ? res.base : null; + } catch { + // searches still work via a literal `~`, only `~` abbreviation degrades + this._serverHome = null; } - this.persistDisabledTools(); + return this._serverHome; } setToolEnabled(key: string, enabled: boolean): void { @@ -472,18 +402,6 @@ class ToolsStore { } } - /** Enable all tools belonging to a specific MCP server */ - enableAllToolsForServer(serverId: string): void { - const connection = mcpStore.getConnections().get(serverId); - - if (!connection) return; - - for (const tool of connection.tools) { - this._disabledTools.delete(this.toolKey(ToolSource.MCP, tool.name, serverId)); - } - this.persistDisabledTools(); - } - toggleGroup(group: ToolGroup): void { const allEnabled = group.tools.every((t) => this.isToolEnabled(t.key)); const target = !allEnabled; @@ -495,8 +413,23 @@ class ToolsStore { this.persistDisabledTools(); } - isGroupFullyEnabled(group: ToolGroup): boolean { - return group.tools.length > 0 && group.tools.every((t) => this.isToolEnabled(t.key)); + toggleTool(key: string): void { + if (this._disabledTools.has(key)) { + this._disabledTools.delete(key); + } else { + this._disabledTools.add(key); + } + + this.persistDisabledTools(); + } + + /** First canonical entry matching a tool name, runtime tool calls resolve by name */ + private findEntryByName(toolName: string): ToolEntry | null { + for (const entry of this.allTools) { + if (entry.definition.function.name === toolName) return entry; + } + + return null; } /** Get MCP tools from health check data, used when live connections aren't established yet */ @@ -524,118 +457,194 @@ class ToolsStore { return result; } - /** First canonical entry matching a tool name, runtime tool calls resolve by name */ - private findEntryByName(toolName: string): ToolEntry | null { - for (const entry of this.allTools) { - if (entry.definition.function.name === toolName) return entry; + private groupLabel(entry: ToolEntry): string { + switch (entry.source) { + case ToolSource.MCP: + return entry.serverName ?? ''; + case ToolSource.CUSTOM: + return TOOL_GROUP_LABELS[ToolSource.CUSTOM]; + case ToolSource.BROWSER: + return TOOL_GROUP_LABELS[ToolSource.BROWSER]; + default: + return TOOL_GROUP_LABELS[ToolSource.SERVER]; } - - return null; } - /** Determine the source of a tool by its name */ - getToolSource(toolName: string): ToolSource | null { - return this.findEntryByName(toolName)?.source ?? null; + private hasServerTool(name: BuiltInTool): boolean { + return this._serverTools.some((def) => def.function.name === name); } - /** Get the display label for the server that owns a given tool */ - getToolServerLabel(toolName: string): string { - const entry = this.findEntryByName(toolName); - - if (!entry) return ''; + private inferTypeFromDefault(value: unknown): string | undefined { + if (typeof value === 'string') return 'string'; - if (entry.serverName) return mcpStore.getServerDisplayName(entry.serverName); + if (typeof value === 'boolean') return 'boolean'; - if (entry.source === ToolSource.SERVER) return TOOL_SERVER_LABELS[ToolSource.SERVER]; + if (typeof value === 'number') return Number.isInteger(value) ? 'integer' : 'number'; - if (entry.source === ToolSource.CUSTOM) return TOOL_SERVER_LABELS[ToolSource.CUSTOM]; + if (Array.isArray(value)) return 'array'; - if (entry.source === ToolSource.BROWSER) return TOOL_SERVER_LABELS[ToolSource.BROWSER]; + if (value !== null && typeof value === 'object') return 'object'; - return ''; + return undefined; } - /** Permission key for a tool name, identical to the selection key */ - getPermissionKey(toolName: string): string | null { - return this.findEntryByName(toolName)?.key ?? null; + private mcpDefinition( + name: string, + description: string | undefined, + schema?: Record + ): OpenAIToolDefinition { + return { + function: { + description, + name, + parameters: schema ?? { properties: {}, required: [], type: JsonSchemaType.OBJECT } + }, + type: ToolCallType.FUNCTION + }; } - /** Check if there are any enabled tools available (server, MCP, or custom) */ - get hasEnabledTools(): boolean { - return this.getEnabledToolsForLLM().length > 0; + /** Normalize MCP tools from live connections when available, fall back to health check data */ + private mcpEntries(): { + serverId: string; + serverName: string; + definition: OpenAIToolDefinition; + }[] { + const out: { serverId: string; serverName: string; definition: OpenAIToolDefinition }[] = []; + const connections = mcpStore.getConnections(); + + if (connections.size > 0) { + for (const [serverId, connection] of connections) { + const serverName = mcpStore.getServerDisplayName(serverId); + + for (const tool of connection.tools) { + const rawSchema = (tool.inputSchema as Record) ?? { + properties: {}, + required: [], + type: JsonSchemaType.OBJECT + }; + + out.push({ + definition: { + function: { + description: tool.description, + name: tool.name, + parameters: this.normalizeJsonSchema(rawSchema) + }, + type: ToolCallType.FUNCTION + }, + serverId, + serverName + }); + } + } + } else { + for (const { serverId, serverName, tools } of this.getMcpToolsFromHealthChecks()) { + for (const tool of tools) { + out.push({ + definition: this.mcpDefinition(tool.name, tool.description), + serverId, + serverName + }); + } + } + } + + return out; } /** - * Check if a working directory is worth setting: at least one server tool - * that reads it is both served and left enabled by the user. + * Recursively normalize a JSON Schema object: infers `type` from `default` + * for properties / items that omit it, and descends into nested `properties` + * and `items`. Returns a new object -- does not mutate the input. */ - get hasEnabledCwdTools(): boolean { - return this._serverTools.some((def) => { - const name = def.function.name; + private normalizeJsonSchema(schema: Record): Record { + if (!schema || typeof schema !== 'object') return schema; - return ( - this._cwdAwareTools.has(name) && - !this._disabledTools.has(this.toolKey(ToolSource.SERVER, name)) - ); - }); - } + const normalized: Record = { ...schema }; - async fetchServerTools(): Promise { - if (this._loading) return; + if (normalized.properties && typeof normalized.properties === 'object') { + const props = normalized.properties as Record>; + const normalizedProps: Record> = {}; - this._loading = true; - this._error = null; - this._toolsEndpointUnreachable = false; + for (const [key, prop] of Object.entries(props)) { + if (!prop || typeof prop !== 'object') { + normalizedProps[key] = prop; - try { - const toolInfos = await ToolsService.list(); + continue; + } - this._serverTools = toolInfos.map((info) => info.definition); - this._cwdAwareTools = new SvelteSet( - toolInfos.filter((info) => info.uses_cwd).map((info) => info.tool) - ); - } catch (err) { - const errorMessage = err instanceof Error ? err.message : String(err); + const normalizedProp: Record = { ...prop }; - this._error = errorMessage; + if (!normalizedProp.type && normalizedProp.default !== undefined) { + const inferred = this.inferTypeFromDefault(normalizedProp.default); - // 403 from /tools means the server was started without --tools - // TODO: check status code instead of relying on message - if (errorMessage.includes('this feature is disabled')) { - this._toolsEndpointUnreachable = true; - console.info('[ToolsStore] Server tools are disabled on the server'); - } else { - console.error('[ToolsStore] Failed to fetch server tools:', err); + if (inferred) normalizedProp.type = inferred; + } + + if (normalizedProp.properties) { + Object.assign( + normalizedProp, + this.normalizeJsonSchema(normalizedProp as Record) + ); + } + + if (normalizedProp.items && typeof normalizedProp.items === 'object') { + normalizedProp.items = this.normalizeJsonSchema( + normalizedProp.items as Record + ); + } + + normalizedProps[key] = normalizedProp; } - } finally { - this._loading = false; + normalized.properties = normalizedProps; + } + + return normalized; + } + + private persistDisabledTools(): void { + try { + localStorage.setItem( + DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY, + JSON.stringify([...this._disabledTools]) + ); + } catch { + // ignore storage errors } } /** - * Absolute home directory on the server, resolved once per session via - * file_glob_search's `base` field (the server expands `~`). Anchors the - * directory picker's search scope and the `~` abbreviation of cwd - * displays. Returns null when tools are unavailable. + * `read_media` runs in the browser on top of the server's `read_file`, so it + * exists only when that tool is served and the active model can perceive the + * bytes. The server cannot make this call - it does not know which model the + * conversation uses. */ - async resolveServerHome(): Promise { - if (this._serverHome !== undefined) return this._serverHome; + private readMediaTool(): OpenAIToolDefinition | null { + if (!this.hasServerTool(BuiltInTool.SERVER_READ_FILE)) return null; - try { - const res = await ToolsService.executeToolRaw(BuiltInTool.SERVER_FILE_GLOB_SEARCH, { - limit: 1, - max_depth: 1, - path: HOME_TILDE, - type: GlobSearchType.DIR - }); + const model = modelsStore.selectedModelName ?? modelsStore.models[0]?.model ?? ''; - this._serverHome = typeof res.base === 'string' ? res.base : null; - } catch { - // searches still work via a literal `~`, only `~` abbreviation degrades - this._serverHome = null; - } + if (!model) return null; - return this._serverHome; + const vision = modelsStore.props.modelSupportsVision(model); + const audio = modelsStore.props.modelSupportsAudio(model); + + if (!vision && !audio) return null; + + return buildReadMediaToolDefinition(vision, audio); + } + + private toolKey(source: ToolSource, name: string, serverId?: string): string { + switch (source) { + case ToolSource.MCP: + return serverId ? `mcp-${serverId}:${name}` : `mcp:${name}`; + case ToolSource.CUSTOM: + return `custom:${name}`; + case ToolSource.BROWSER: + return `browser:${name}`; + default: + return `server:${name}`; + } } } diff --git a/tools/ui/src/lib/types/agentic.d.ts b/tools/ui/src/lib/types/agentic.d.ts index e7c6d34e1dc5..1a604476bf9d 100644 --- a/tools/ui/src/lib/types/agentic.d.ts +++ b/tools/ui/src/lib/types/agentic.d.ts @@ -205,7 +205,7 @@ export interface AgenticSection { /** ID of the model-side tool call (matches tool_calls[i].id). Lets * downstream consumers correlate a section with the agentic loop's * currently-executing tool, e.g. to drive live-streaming UI state - * by matching against agenticStore.executingToolCallId. */ + * by matching against agenticStore.getExecutingToolCallId. */ toolCallId?: string; wasInterrupted?: boolean; } diff --git a/tools/ui/src/lib/utils/api-fetch.ts b/tools/ui/src/lib/utils/api-fetch.ts index 65e1129def12..20592000493b 100644 --- a/tools/ui/src/lib/utils/api-fetch.ts +++ b/tools/ui/src/lib/utils/api-fetch.ts @@ -1,7 +1,6 @@ import { getAuthHeaders, getJsonHeaders } from './api-headers'; import { base } from '$app/paths'; -import { ERROR_MESSAGES, HTTP_CODE_TO_STRING } from '$lib/constants'; -import { UrlProtocol } from '$lib/enums'; +import { API_ABSOLUTE_URL_PROTOCOLS, ERROR_MESSAGES, HTTP_CODE_TO_STRING } from '$lib/constants'; /** * API Fetch Utilities @@ -63,10 +62,8 @@ export async function apiFetch(path: string, options: ApiFetchOptions = {}): const { authOnly = false, headers: customHeaders, ...fetchOptions } = options; const baseHeaders = authOnly ? getAuthHeaders() : getJsonHeaders(); const headers = { ...baseHeaders, ...customHeaders }; - const url = - path.startsWith(UrlProtocol.HTTP) || path.startsWith(UrlProtocol.HTTPS) - ? path - : `${base}${path}`; + // absolute URLs with an allowed protocol pass through untouched; relative paths get the base prefix + const url = API_ABSOLUTE_URL_PROTOCOLS.some((p) => path.startsWith(p)) ? path : `${base}${path}`; let response; @@ -117,28 +114,7 @@ export async function apiFetchWithParams( } } - const { authOnly = false, headers: customHeaders, ...fetchOptions } = options; - const baseHeaders = authOnly ? getAuthHeaders() : getJsonHeaders(); - const headers = { ...baseHeaders, ...customHeaders }; - - let response; - - try { - response = await fetch(url.toString(), { - ...fetchOptions, - headers - }); - } catch (e) { - throw new Error(beautifyNetworkError(e)); - } - - if (!response.ok) { - const errorMessage = await parseErrorMessage(response); - - throw new ApiError(errorMessage, response.status); - } - - return response.json() as Promise; + return apiFetch(url.toString(), options); } /** diff --git a/tools/ui/src/lib/utils/api-headers.ts b/tools/ui/src/lib/utils/api-headers.ts index 4b2b19d442e6..49d56d06192c 100644 --- a/tools/ui/src/lib/utils/api-headers.ts +++ b/tools/ui/src/lib/utils/api-headers.ts @@ -1,7 +1,7 @@ import { redactValue } from './redact'; import { CORS_PROXY, HEADERS } from '$lib/constants'; import { MimeTypeApplication } from '$lib/enums'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; /** * Get authorization headers for API requests diff --git a/tools/ui/src/lib/utils/api-key-validation.ts b/tools/ui/src/lib/utils/api-key-validation.ts index 8cde154fd2f9..187199afc262 100644 --- a/tools/ui/src/lib/utils/api-key-validation.ts +++ b/tools/ui/src/lib/utils/api-key-validation.ts @@ -3,7 +3,7 @@ import { browser } from '$app/environment'; import { base } from '$app/paths'; import { HEADERS } from '$lib/constants'; import { MimeTypeApplication } from '$lib/enums'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; /** * Validates API key by making a request to the server props endpoint diff --git a/tools/ui/src/lib/utils/audio-recording.ts b/tools/ui/src/lib/utils/audio-recording.ts index 1241d6e05593..4cfe17378480 100644 --- a/tools/ui/src/lib/utils/audio-recording.ts +++ b/tools/ui/src/lib/utils/audio-recording.ts @@ -14,10 +14,37 @@ import { MimeTypeAudio } from '$lib/enums'; * - Proper cleanup and resource management */ export class AudioRecorder { - private mediaRecorder: MediaRecorder | null = null; private audioChunks: Blob[] = []; - private stream: MediaStream | null = null; + private mediaRecorder: MediaRecorder | null = null; private recordingState: boolean = false; + private stream: MediaStream | null = null; + + cancelRecording(): void { + const recorder = this.mediaRecorder; + const stream = this.stream; + + this.mediaRecorder = null; + this.audioChunks = []; + this.stream = null; + this.recordingState = false; + + if (recorder && recorder.state !== 'inactive') { + // Drop the original handlers so the pending stop event does not touch the instance + recorder.onstop = null; + recorder.onerror = null; + recorder.stop(); + } + + if (stream) { + for (const track of stream.getTracks()) { + track.stop(); + } + } + } + + isRecording(): boolean { + return this.recordingState; + } async startRecording(): Promise { try { @@ -90,33 +117,6 @@ export class AudioRecorder { }); } - isRecording(): boolean { - return this.recordingState; - } - - cancelRecording(): void { - const recorder = this.mediaRecorder; - const stream = this.stream; - - this.mediaRecorder = null; - this.audioChunks = []; - this.stream = null; - this.recordingState = false; - - if (recorder && recorder.state !== 'inactive') { - // Drop the original handlers so the pending stop event does not touch the instance - recorder.onstop = null; - recorder.onerror = null; - recorder.stop(); - } - - if (stream) { - for (const track of stream.getTracks()) { - track.stop(); - } - } - } - private initializeRecorder(stream: MediaStream): void { const options: MediaRecorderOptions = {}; diff --git a/tools/ui/src/lib/utils/cache-ttl.ts b/tools/ui/src/lib/utils/cache-ttl.ts index bb0100755b30..bec40989c4c1 100644 --- a/tools/ui/src/lib/utils/cache-ttl.ts +++ b/tools/ui/src/lib/utils/cache-ttl.ts @@ -31,9 +31,29 @@ interface CacheEntry { export class TTLCache { private cache = new Map>(); - private readonly ttlMs: number; private readonly maxEntries: number; private readonly onEvict?: (key: string, value: unknown) => void; + private readonly ttlMs: number; + + /** + * Get the number of entries (including potentially expired ones). + */ + get size(): number { + return this.cache.size; + } + + /** + * Clear all entries from cache. + */ + clear(): void { + if (this.onEvict) { + for (const [key, entry] of this.cache) { + this.onEvict(key, entry.value); + } + } + + this.cache.clear(); + } constructor(options: TTLCacheOptions = {}) { this.ttlMs = options.ttlMs ?? CACHE.DEFAULT_TTL_MS; @@ -41,6 +61,19 @@ export class TTLCache { this.onEvict = options.onEvict; } + /** + * Delete a specific key from cache. + */ + delete(key: K): boolean { + const entry = this.cache.get(key); + + if (entry && this.onEvict) { + this.onEvict(key, entry.value); + } + + return this.cache.delete(key); + } + /** * Get a value from cache. Returns null if expired or not found. */ @@ -61,25 +94,6 @@ export class TTLCache { return entry.value; } - /** - * Set a value in cache with TTL. - */ - set(key: K, value: V, customTtlMs?: number): void { - // Evict oldest entries if at capacity - if (this.cache.size >= this.maxEntries && !this.cache.has(key)) { - this.evictOldest(); - } - - const ttl = customTtlMs ?? this.ttlMs; - const now = Date.now(); - - this.cache.set(key, { - expiresAt: now + ttl, - lastAccessed: now, - value - }); - } - /** * Check if key exists and is not expired. */ @@ -98,36 +112,19 @@ export class TTLCache { } /** - * Delete a specific key from cache. + * Get all valid (non-expired) keys. */ - delete(key: K): boolean { - const entry = this.cache.get(key); - - if (entry && this.onEvict) { - this.onEvict(key, entry.value); - } - - return this.cache.delete(key); - } + keys(): K[] { + const now = Date.now(); + const validKeys: K[] = []; - /** - * Clear all entries from cache. - */ - clear(): void { - if (this.onEvict) { - for (const [key, entry] of this.cache) { - this.onEvict(key, entry.value); + for (const [key, entry] of this.cache) { + if (now <= entry.expiresAt) { + validKeys.push(key); } } - this.cache.clear(); - } - - /** - * Get the number of entries (including potentially expired ones). - */ - get size(): number { - return this.cache.size; + return validKeys; } /** @@ -150,38 +147,22 @@ export class TTLCache { } /** - * Get all valid (non-expired) keys. + * Set a value in cache with TTL. */ - keys(): K[] { - const now = Date.now(); - const validKeys: K[] = []; - - for (const [key, entry] of this.cache) { - if (now <= entry.expiresAt) { - validKeys.push(key); - } + set(key: K, value: V, customTtlMs?: number): void { + // Evict oldest entries if at capacity + if (this.cache.size >= this.maxEntries && !this.cache.has(key)) { + this.evictOldest(); } - return validKeys; - } - - /** - * Evict the oldest (least recently accessed) entry. - */ - private evictOldest(): void { - let oldestKey: K | null = null; - let oldestTime = Infinity; - - for (const [key, entry] of this.cache) { - if (entry.lastAccessed < oldestTime) { - oldestTime = entry.lastAccessed; - oldestKey = key; - } - } + const ttl = customTtlMs ?? this.ttlMs; + const now = Date.now(); - if (oldestKey !== null) { - this.delete(oldestKey); - } + this.cache.set(key, { + expiresAt: now + ttl, + lastAccessed: now, + value + }); } /** @@ -205,6 +186,25 @@ export class TTLCache { return true; } + + /** + * Evict the oldest (least recently accessed) entry. + */ + private evictOldest(): void { + let oldestKey: K | null = null; + let oldestTime = Infinity; + + for (const [key, entry] of this.cache) { + if (entry.lastAccessed < oldestTime) { + oldestTime = entry.lastAccessed; + oldestKey = key; + } + } + + if (oldestKey !== null) { + this.delete(oldestKey); + } + } } /** @@ -213,14 +213,26 @@ export class TTLCache { */ export class ReactiveTTLMap { private entries = $state>>(new Map()); - private readonly ttlMs: number; private readonly maxEntries: number; + private readonly ttlMs: number; + + get size(): number { + return this.entries.size; + } + + clear(): void { + this.entries.clear(); + } constructor(options: TTLCacheOptions = {}) { this.ttlMs = options.ttlMs ?? CACHE.DEFAULT_TTL_MS; this.maxEntries = options.maxEntries ?? CACHE.DEFAULT_MAX_ENTRIES; } + delete(key: K): boolean { + return this.entries.delete(key); + } + get(key: K): V | null { const entry = this.entries.get(key); @@ -237,21 +249,6 @@ export class ReactiveTTLMap { return entry.value; } - set(key: K, value: V, customTtlMs?: number): void { - if (this.entries.size >= this.maxEntries && !this.entries.has(key)) { - this.evictOldest(); - } - - const ttl = customTtlMs ?? this.ttlMs; - const now = Date.now(); - - this.entries.set(key, { - expiresAt: now + ttl, - lastAccessed: now, - value - }); - } - has(key: K): boolean { const entry = this.entries.get(key); @@ -266,18 +263,6 @@ export class ReactiveTTLMap { return true; } - delete(key: K): boolean { - return this.entries.delete(key); - } - - clear(): void { - this.entries.clear(); - } - - get size(): number { - return this.entries.size; - } - prune(): number { const now = Date.now(); @@ -293,6 +278,21 @@ export class ReactiveTTLMap { return pruned; } + set(key: K, value: V, customTtlMs?: number): void { + if (this.entries.size >= this.maxEntries && !this.entries.has(key)) { + this.evictOldest(); + } + + const ttl = customTtlMs ?? this.ttlMs; + const now = Date.now(); + + this.entries.set(key, { + expiresAt: now + ttl, + lastAccessed: now, + value + }); + } + private evictOldest(): void { let oldestKey: K | null = null; let oldestTime = Infinity; diff --git a/tools/ui/src/lib/utils/chat-form-input-rich-tokenizer.ts b/tools/ui/src/lib/utils/chat-form-input-rich-tokenizer.ts index c09afd018bcc..626b10b29b52 100644 --- a/tools/ui/src/lib/utils/chat-form-input-rich-tokenizer.ts +++ b/tools/ui/src/lib/utils/chat-form-input-rich-tokenizer.ts @@ -38,7 +38,7 @@ import { SETTINGS_KEYS } from '$lib/constants'; import { BooleanString, ChatFormInputRichTokenKind } from '$lib/enums'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { toolsStore } from '$lib/stores/tools.svelte'; import type { ChatFormInputRichToken } from '$lib/types/chat-form-input-rich'; diff --git a/tools/ui/src/lib/utils/convert-files-to-extra.ts b/tools/ui/src/lib/utils/convert-files-to-extra.ts index e348f25fe9c3..735e91c44a4b 100644 --- a/tools/ui/src/lib/utils/convert-files-to-extra.ts +++ b/tools/ui/src/lib/utils/convert-files-to-extra.ts @@ -4,8 +4,8 @@ import { isLikelyTextFile, readFileAsText } from './text-files'; import { isWebpMimeType, webpBase64UrlToPngDataURL } from './webp-to-png'; import { SETTINGS_KEYS } from '$lib/constants'; import { AttachmentType, FileTypeCategory, SpecialFileType } from '$lib/enums'; -import { modelsStore } from '$lib/stores/models.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { ChatUploadedFile, DatabaseMessageExtra, FileProcessingResult } from '$lib/types'; import { getFileTypeCategory } from '$lib/utils'; import { toast } from 'svelte-sonner'; @@ -112,7 +112,7 @@ export async function parseFilesToMessageExtras( const currentConfig = settingsStore.config; // Use per-model vision check for router mode const hasVisionSupport = activeModelId - ? modelsStore.modelSupportsVision(activeModelId) + ? modelsStore.props.modelSupportsVision(activeModelId) : false; // Force PDF-to-text for non-vision models diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts index 239f9f5724da..079cdc871c60 100644 --- a/tools/ui/src/lib/utils/index.ts +++ b/tools/ui/src/lib/utils/index.ts @@ -130,7 +130,7 @@ export { getImageErrorFallbackHtml } from './image-error-fallback'; // SSE-with-JSON stream iterator (used by server tool streaming, decoupled // from chat.service.ts which embeds its own SSE parser for resume support) -export { parseSseJsonStream } from './sse'; +export { extractSseDataPayload, parseSseJsonStream, splitSseRecords } from './sse'; // Stream session identity (conversation-id based) export { streamIdentity } from './stream-identity'; @@ -150,7 +150,10 @@ export { getResourceIcon, getResourceTextContent, getResourceBlobContent, - downloadResourceContent + downloadResourceContent, + getMcpIconUrl, + getMcpServerFaviconFallback, + getMcpServerLabel } from './mcp'; // URI Template utilities diff --git a/tools/ui/src/lib/utils/mcp.ts b/tools/ui/src/lib/utils/mcp.ts index 61d5f8a9a508..c60a59e80e3d 100644 --- a/tools/ui/src/lib/utils/mcp.ts +++ b/tools/ui/src/lib/utils/mcp.ts @@ -1,3 +1,4 @@ +import { extractRootDomain } from './url'; import { AlertTriangle, Code, @@ -12,8 +13,10 @@ import { CODE_FILE_EXTENSION_REGEX, DEFAULT_RESOURCE_FILENAME, DISPLAY_NAME_SEPARATOR_REGEX, + EXPECTED_THEMED_ICON_PAIR_COUNT, FILE_EXTENSION_REGEX, IMAGE_FILE_EXTENSION_REGEX, + MCP_ALLOWED_ICON_MIME_TYPES, MCP_SERVER_ID_PREFIX, MCP_SSE, MIME_TYPE_PREFIXES, @@ -24,8 +27,22 @@ import { TEXT_FILE_EXTENSION_REGEX, URI_PATTERNS } from '$lib/constants'; -import { MCPLogLevel, MCPTransportType, MimeTypeText, UrlProtocol } from '$lib/enums'; -import type { MCPResourceContent, MCPResourceInfo, MCPServerSettingsEntry } from '$lib/types'; +import { + ColorMode, + HealthCheckStatus, + MCPLogLevel, + MCPTransportType, + MimeTypeText, + UrlProtocol +} from '$lib/enums'; +import type { + HealthCheckState, + MCPResourceContent, + MCPResourceIcon, + MCPResourceInfo, + MCPServerDisplayInfo, + MCPServerSettingsEntry +} from '$lib/types'; import type { MimeTypeUnion } from '$lib/types/common'; import type { Component } from 'svelte'; @@ -316,3 +333,132 @@ export function downloadResourceContent( document.body.removeChild(a); URL.revokeObjectURL(url); } + +/** + * Validates that an icon URI uses a safe scheme (https: or data:). + */ +function isValidMcpIconUri(src: string): boolean { + try { + if (src.startsWith(UrlProtocol.DATA)) return true; + + const url = new URL(src); + + return url.protocol === UrlProtocol.HTTPS; + } catch { + return false; + } +} + +/** + * Selects the best icon URL from an MCP icons array. + * Follows security guidelines from the MCP specification: + * - Only allows https: and data: URIs + * - Filters to supported MIME types + * + * Selection priority: + * 1. Icon matching the current color scheme (dark/light) + * 2. Universal icon (no theme specified); if exactly 2, assumes [0]=light, [1]=dark + * 3. First valid icon as last resort + */ +export function getMcpIconUrl(icons: MCPResourceIcon[] | undefined, isDark = false): string | null { + if (!icons?.length) return null; + + const validIcons = icons.filter((icon) => { + if (!icon.src || !isValidMcpIconUri(icon.src)) return false; + + if (icon.mimeType && !MCP_ALLOWED_ICON_MIME_TYPES.has(icon.mimeType)) return false; + + return true; + }); + + if (validIcons.length === 0) return null; + + const preferredTheme = isDark ? ColorMode.DARK : ColorMode.LIGHT; + // 1. Prefer icon explicitly matching the current color scheme + const themedIcon = validIcons.find((icon) => icon.theme === preferredTheme); + + if (themedIcon) return themedIcon.src; + + // 2. Handle universal icons (no theme specified) + const universalIcons = validIcons.filter((icon) => !icon.theme); + + if (universalIcons.length === EXPECTED_THEMED_ICON_PAIR_COUNT) { + // Heuristic: two theme-less icons → assume [0] = light, [1] = dark + return universalIcons[isDark ? 1 : 0].src; + } + + if (universalIcons.length > 0) { + return universalIcons[0].src; + } + + // 3. Last resort: use opposite-theme icon + return validIcons[0].src; +} + +/** + * Construct a fallback favicon URL from the MCP server URL. + * e.g. https://mcp.example.com/sse -> https://example.com/favicon.ico + */ +export function getMcpServerFaviconFallback(serverUrl: string): string | null { + try { + const url = new URL(serverUrl); + const rootDomain = extractRootDomain(url); + + if (!rootDomain) return null; + + const origin = `${url.protocol}//${rootDomain}`; + const candidates = ['favicon.ico', 'favicon.png']; + + for (const path of candidates) { + const faviconUrl = `${origin}/${path}`; + + if (isValidMcpIconUri(faviconUrl)) { + return faviconUrl; + } + } + } catch { + // Invalid URL, return null + } + + return null; +} + +/** + * Resolves the raw label for a server: user-defined display name first, + * then server-reported title or name when the health check succeeded, + * then the configured name (admin baseline or legacy data), then URL. + */ +function getMcpServerBaseLabel( + server: MCPServerDisplayInfo, + healthState?: HealthCheckState +): string { + if (server.displayName) return server.displayName; + + if (healthState?.status === HealthCheckStatus.SUCCESS) + return ( + healthState.serverInfo?.title || healthState.serverInfo?.name || server.name || server.url + ); + + return server.name || server.url; +} + +/** + * Returns the display label for a server, suffixed with a positional + * counter when several configured servers resolve to the same base label + * (e.g. two endpoints of the same host reporting an identical name). + * Numbering follows config order, so it is stable across renders. + */ +export function getMcpServerLabel( + server: MCPServerDisplayInfo, + servers: MCPServerDisplayInfo[], + healthChecks: Record +): string { + const label = getMcpServerBaseLabel(server, healthChecks[server.id]); + const twins = servers.filter((s) => getMcpServerBaseLabel(s, healthChecks[s.id]) === label); + + if (twins.length < 2) return label; + + const position = twins.findIndex((s) => s.id === server.id); + + return position < 0 ? label : `${label} (${position + 1})`; +} diff --git a/tools/ui/src/lib/utils/process-uploaded-files.ts b/tools/ui/src/lib/utils/process-uploaded-files.ts index 49bdd2412fa5..e71371345c24 100644 --- a/tools/ui/src/lib/utils/process-uploaded-files.ts +++ b/tools/ui/src/lib/utils/process-uploaded-files.ts @@ -4,8 +4,8 @@ import { isSvgMimeType, svgBase64UrlToPngDataURL } from './svg-to-png'; import { isWebpMimeType, webpBase64UrlToPngDataURL } from './webp-to-png'; import { SETTINGS_KEYS } from '$lib/constants'; import { FileTypeCategory } from '$lib/enums'; -import { modelsStore } from '$lib/stores/models.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { getFileTypeCategory } from '$lib/utils'; import { toast } from 'svelte-sonner'; @@ -108,7 +108,7 @@ export async function processFilesToChatUploaded( // Show suggestion toast if vision model is available but PDF as image is disabled const hasVisionSupport = activeModelId - ? modelsStore.modelSupportsVision(activeModelId) + ? modelsStore.props.modelSupportsVision(activeModelId) : false; const currentConfig = settingsStore.config; diff --git a/tools/ui/src/lib/utils/source-history.ts b/tools/ui/src/lib/utils/source-history.ts index 32995ae03465..6228ae7e49dd 100644 --- a/tools/ui/src/lib/utils/source-history.ts +++ b/tools/ui/src/lib/utils/source-history.ts @@ -12,9 +12,9 @@ export interface SourceHistoryEntry { } export class SourceHistory { - private undoStack: SourceHistoryEntry[] = []; - private redoStack: SourceHistoryEntry[] = []; private lastPush = 0; + private redoStack: SourceHistoryEntry[] = []; + private undoStack: SourceHistoryEntry[] = []; constructor( private limit = 100, @@ -32,24 +32,24 @@ export class SourceHistory { this.redoStack = []; } - undo(current: SourceHistoryEntry): SourceHistoryEntry | null { - const entry = this.undoStack.pop(); + redo(current: SourceHistoryEntry): SourceHistoryEntry | null { + const entry = this.redoStack.pop(); if (!entry) return null; - this.redoStack.push(current); - this.lastPush = 0; // the next edit after an undo starts a new group + this.undoStack.push(current); + this.lastPush = 0; return entry; } - redo(current: SourceHistoryEntry): SourceHistoryEntry | null { - const entry = this.redoStack.pop(); + undo(current: SourceHistoryEntry): SourceHistoryEntry | null { + const entry = this.undoStack.pop(); if (!entry) return null; - this.undoStack.push(current); - this.lastPush = 0; + this.redoStack.push(current); + this.lastPush = 0; // the next edit after an undo starts a new group return entry; } diff --git a/tools/ui/src/lib/utils/sse.ts b/tools/ui/src/lib/utils/sse.ts index 41d9a1152a01..c984e77ee67d 100644 --- a/tools/ui/src/lib/utils/sse.ts +++ b/tools/ui/src/lib/utils/sse.ts @@ -25,6 +25,30 @@ export interface SseJsonEvent { data: T; } +/** + * Splits a raw SSE byte buffer into complete records on the blank-line + * boundary, returning the leftover partial record separately. Shared by the + * record-based consumers (parseSseJsonStream, models.service). + */ +export function splitSseRecords(buffer: string): { records: string[]; rest: string } { + const parts = buffer.split(SSE_RECORD_SEPARATOR); + + return { records: parts.slice(0, -1), rest: parts[parts.length - 1] ?? '' }; +} + +/** + * Extracts the joined `data:` payload from one SSE record (the data lines + * concatenated with a newline), or an empty string when the record carries + * no data lines. Used by models.service to parse status envelopes. + */ +export function extractSseDataPayload(record: string): string { + return record + .split(SSE_LINE_SEPARATOR) + .filter((line) => line.startsWith(SSE_DATA_PREFIX)) + .map((line) => line.slice(SSE_DATA_PREFIX.length).trim()) + .join(SSE_LINE_SEPARATOR); +} + export async function* parseSseJsonStream( response: Response, signal?: AbortSignal @@ -46,9 +70,9 @@ export async function* parseSseJsonStream( if (done) break; buffer += decoder.decode(value, { stream: true }); - const records = buffer.split(SSE_RECORD_SEPARATOR); + const { records, rest } = splitSseRecords(buffer); - buffer = records.pop() ?? ''; + buffer = rest; for (const record of records) { if (!record) continue; diff --git a/tools/ui/src/routes/(chat)/+page.svelte b/tools/ui/src/routes/(chat)/+page.svelte index 224d264c4391..de8574e3528c 100644 --- a/tools/ui/src/routes/(chat)/+page.svelte +++ b/tools/ui/src/routes/(chat)/+page.svelte @@ -47,8 +47,8 @@ serverStore.isRouterMode && !modelsStore.isModelLoaded(model.id) ) { - modelsStore - .loadModel(model.id) + modelsStore.status + .load(model.id) .catch((error) => console.error('Failed to load model:', error)); } } catch (error) { @@ -77,7 +77,7 @@ onMount(async () => { if (!conversationsStore.isInitialized) { - await conversationsStore.init(); + await conversationsStore.initialize(); } conversationsStore.clearActiveConversation(); diff --git a/tools/ui/src/routes/+layout.svelte b/tools/ui/src/routes/+layout.svelte index 8314cd2a2dbe..f87bbe26a235 100644 --- a/tools/ui/src/routes/+layout.svelte +++ b/tools/ui/src/routes/+layout.svelte @@ -216,11 +216,11 @@ if (!serverStore.isRouterMode) return; untrack(() => { - modelsStore.subscribeStatus(); + modelsStore.status.subscribe(); }); return () => { - modelsStore.unsubscribeStatus(); + modelsStore.status.unsubscribe(); }; }); diff --git a/tools/ui/tests/client/agentic-stream.perf.svelte.test.ts b/tools/ui/tests/client/agentic-stream.perf.svelte.test.ts index 0b06d57a5b41..b4d6df453889 100644 --- a/tools/ui/tests/client/agentic-stream.perf.svelte.test.ts +++ b/tools/ui/tests/client/agentic-stream.perf.svelte.test.ts @@ -14,7 +14,7 @@ import { perfState } from './components/agentic-perf-state.svelte'; import AgenticPerfWrapper from './components/AgenticPerfWrapper.svelte'; import ChatMessagesPerfWrapper from './components/ChatMessagesPerfWrapper.svelte'; import { MessageRole } from '$lib/enums'; -import { conversationsStore } from '$lib/stores/conversations.svelte'; +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; import type { DatabaseMessage } from '$lib/types'; import { tick } from 'svelte'; import { describe, it } from 'vitest'; diff --git a/tools/ui/tests/client/apikey-splash.svelte.test.ts b/tools/ui/tests/client/apikey-splash.svelte.test.ts index bad7f6ccb076..b2705dd8ca24 100644 --- a/tools/ui/tests/client/apikey-splash.svelte.test.ts +++ b/tools/ui/tests/client/apikey-splash.svelte.test.ts @@ -1,5 +1,5 @@ import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { validateApiKey } from '$lib/utils/api-key-validation'; import { beforeEach, describe, expect, it } from 'vitest'; diff --git a/tools/ui/tests/client/chat-form-enter-code-block.svelte.test.ts b/tools/ui/tests/client/chat-form-enter-code-block.svelte.test.ts index 485dc39655a6..3454170b6b4c 100644 --- a/tools/ui/tests/client/chat-form-enter-code-block.svelte.test.ts +++ b/tools/ui/tests/client/chat-form-enter-code-block.svelte.test.ts @@ -7,7 +7,7 @@ import ChatFormTestWrapper from './components/ChatFormTestWrapper.svelte'; import { SETTINGS_KEYS } from '$lib/constants'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { tick } from 'svelte'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { userEvent } from 'vitest/browser'; diff --git a/tools/ui/tests/client/components/ChatMessagesPerfWrapper.svelte b/tools/ui/tests/client/components/ChatMessagesPerfWrapper.svelte index 504f685973c9..ab5cc38bc96c 100644 --- a/tools/ui/tests/client/components/ChatMessagesPerfWrapper.svelte +++ b/tools/ui/tests/client/components/ChatMessagesPerfWrapper.svelte @@ -4,7 +4,7 @@ // toolMessages array) rather than a single message subtree. import ChatMessages from '$lib/components/app/chat/ChatMessages/ChatMessages.svelte'; import * as Tooltip from '$lib/components/ui/tooltip'; - import { conversationsStore } from '$lib/stores/conversations.svelte'; + import { conversationsStore } from '$lib/stores/conversations/index.svelte'; diff --git a/tools/ui/tests/client/mcp-display-name.svelte.test.ts b/tools/ui/tests/client/mcp-display-name.svelte.test.ts index f17e08cf1b51..7db0ffd42e5a 100644 --- a/tools/ui/tests/client/mcp-display-name.svelte.test.ts +++ b/tools/ui/tests/client/mcp-display-name.svelte.test.ts @@ -1,6 +1,6 @@ import { McpServerForm } from '$lib/components/app/mcp'; -import { mcpStore } from '$lib/stores/mcp.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { mcpStore } from '$lib/stores/mcp/index.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { beforeEach, describe, expect, it } from 'vitest'; import { render } from 'vitest-browser-svelte'; diff --git a/tools/ui/tests/client/sandbox.service.svelte.test.ts b/tools/ui/tests/client/sandbox.service.svelte.test.ts index 7c0d7926f845..547e3ac1f291 100644 --- a/tools/ui/tests/client/sandbox.service.svelte.test.ts +++ b/tools/ui/tests/client/sandbox.service.svelte.test.ts @@ -10,7 +10,7 @@ const run = (code: string, timeoutMs?: number) => describe('sandbox service', () => { beforeEach(async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.config = { ...settingsStore.config, diff --git a/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts b/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts index 45af7e0d151c..0ed6996b5304 100644 --- a/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts +++ b/tools/ui/tests/client/settings-registry-invariants.svelte.test.ts @@ -1,7 +1,7 @@ import { CONFIG_LOCALSTORAGE_KEY, SETTING_CONFIG_DEFAULT } from '$lib/constants'; import { ParameterSyncService } from '$lib/services/parameter-sync.service'; import { serverStore } from '$lib/stores/server.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import type { SettingsConfigType } from '$lib/types'; import { beforeEach, describe, expect, it } from 'vitest'; diff --git a/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts index 32f4ff3dd4bd..ce65aeb70001 100644 --- a/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts +++ b/tools/ui/tests/client/settings-render-keys-migration.svelte.test.ts @@ -6,7 +6,7 @@ import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; import { MigrationService } from '$lib/services/migration.service'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { beforeEach, describe, expect, it } from 'vitest'; const RENDER_KEYS_MIGRATION_ID = 'render-keys-unfold-v1'; diff --git a/tools/ui/tests/client/ui-settings-sync.svelte.test.ts b/tools/ui/tests/client/ui-settings-sync.svelte.test.ts index 6dca891c85b3..ca9268e2e067 100644 --- a/tools/ui/tests/client/ui-settings-sync.svelte.test.ts +++ b/tools/ui/tests/client/ui-settings-sync.svelte.test.ts @@ -1,6 +1,6 @@ import { CONFIG_LOCALSTORAGE_KEY } from '$lib/constants'; import { serverStore } from '$lib/stores/server.svelte'; -import { settingsStore } from '$lib/stores/settings.svelte'; +import { settingsStore } from '$lib/stores/settings/index.svelte'; import { beforeEach, describe, expect, it } from 'vitest'; function mockProps(uiSettings: Record) { diff --git a/tools/ui/tests/client/update-message-in-place.svelte.test.ts b/tools/ui/tests/client/update-message-in-place.svelte.test.ts index 65298b44b890..ea3b65d0cc26 100644 --- a/tools/ui/tests/client/update-message-in-place.svelte.test.ts +++ b/tools/ui/tests/client/update-message-in-place.svelte.test.ts @@ -8,7 +8,7 @@ // -> 3.07ms at 40). Mutating in place keeps it flat. import { MessageRole } from '$lib/enums'; -import { conversationsStore } from '$lib/stores/conversations.svelte'; +import { conversationsStore } from '$lib/stores/conversations/index.svelte'; import type { DatabaseMessage } from '$lib/types'; import { describe, expect, it } from 'vitest'; diff --git a/tools/ui/tests/stories/ChatMessage.stories.svelte b/tools/ui/tests/stories/ChatMessage.stories.svelte index 84fee2ea1c78..e9bf7a6f6fb0 100644 --- a/tools/ui/tests/stories/ChatMessage.stories.svelte +++ b/tools/ui/tests/stories/ChatMessage.stories.svelte @@ -105,7 +105,7 @@ message: userMessage }} play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', false); }} @@ -118,7 +118,7 @@ message: assistantMessage }} play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', false); }} @@ -131,7 +131,7 @@ message: assistantWithReasoning }} play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', false); }} @@ -144,7 +144,7 @@ message: rawOutputMessage }} play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', true); }} @@ -157,7 +157,7 @@ }} asChild play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', false); // Phase 1: Stream reasoning content in chunks @@ -213,11 +213,11 @@ message: processingMessage }} play={async () => { - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); settingsStore.updateConfig('showRawOutputSwitch', false); // Import the chat store to simulate loading state - const { chatStore } = await import('$lib/stores/chat.svelte'); + const { chatStore } = await import('$lib/stores/chat/index.svelte'); // Set loading state to true to trigger the processing UI chatStore.isLoading = true; diff --git a/tools/ui/tests/stories/ModelsSelector.stories.svelte b/tools/ui/tests/stories/ModelsSelector.stories.svelte index d63300cb211d..7018d09e7b8a 100644 --- a/tools/ui/tests/stories/ModelsSelector.stories.svelte +++ b/tools/ui/tests/stories/ModelsSelector.stories.svelte @@ -4,7 +4,7 @@ import ModelsSelectorOption from '$lib/components/app/models/ModelsSelectorOption.svelte'; import type { GroupedModelOptions, ModelItem } from '$lib/components/app/models/utils'; import { ServerModelStatus } from '$lib/enums'; - import { modelsStore } from '$lib/stores/models.svelte'; + import { modelsStore } from '$lib/stores/models/index.svelte'; const { Story } = defineMeta({ parameters: { diff --git a/tools/ui/tests/stories/SidebarNavigation.stories.svelte b/tools/ui/tests/stories/SidebarNavigation.stories.svelte index 635992601200..ddaa90485d33 100644 --- a/tools/ui/tests/stories/SidebarNavigation.stories.svelte +++ b/tools/ui/tests/stories/SidebarNavigation.stories.svelte @@ -53,7 +53,7 @@ asChild name="Default" play={async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); waitFor(() => setTimeout(() => { @@ -71,7 +71,7 @@ asChild name="SearchActive" play={async ({ userEvent }) => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); waitFor(() => setTimeout(() => { @@ -98,7 +98,7 @@ name="Empty" play={async () => { // Mock empty conversations store - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.conversations = []; }} diff --git a/tools/ui/tests/stories/fixtures/storybook-mocks.ts b/tools/ui/tests/stories/fixtures/storybook-mocks.ts index 736674690409..ac9fb63cd077 100644 --- a/tools/ui/tests/stories/fixtures/storybook-mocks.ts +++ b/tools/ui/tests/stories/fixtures/storybook-mocks.ts @@ -1,4 +1,4 @@ -import { modelsStore } from '$lib/stores/models.svelte'; +import { modelsStore } from '$lib/stores/models/index.svelte'; import { serverStore } from '$lib/stores/server.svelte'; /** diff --git a/tools/ui/tests/unit/chat-activity.test.ts b/tools/ui/tests/unit/chat-activity.test.ts new file mode 100644 index 000000000000..051648ead62f --- /dev/null +++ b/tools/ui/tests/unit/chat-activity.test.ts @@ -0,0 +1,77 @@ +import { ChatActivityStore } from '$lib/stores/chat/activity.svelte'; +import { beforeEach, describe, expect, it } from 'vitest'; + +describe('ChatActivityStore', () => { + let store: ChatActivityStore; + + beforeEach(() => { + store = new ChatActivityStore(); + }); + + it('starts with no local or remote activity', () => { + expect(store.loadingConvs).toEqual([]); + expect(store.isLocal('a')).toBe(false); + expect(store.isRemote('a')).toBe(false); + }); + + it('markLocal adds a conv to the local set and the loading union', () => { + store.markLocal('a'); + + expect(store.isLocal('a')).toBe(true); + expect(store.isRemote('a')).toBe(false); + expect(store.loadingConvs).toEqual(['a']); + }); + + it('localEnded removes a local conv', () => { + store.markLocal('a'); + store.localEnded('a'); + + expect(store.isLocal('a')).toBe(false); + expect(store.loadingConvs).toEqual([]); + }); + + it('localEnded also drops a stale remote hint for the same conv', () => { + store.markLocal('a'); + store.applyRemoteSnapshot(['a']); + expect(store.isRemote('a')).toBe(true); + + store.localEnded('a'); + + expect(store.isLocal('a')).toBe(false); + expect(store.isRemote('a')).toBe(false); + expect(store.loadingConvs).toEqual([]); + }); + + it('applyRemoteSnapshot adds remote convs and unions them with local', () => { + store.markLocal('local'); + store.applyRemoteSnapshot(['remote']); + + expect(store.isRemote('remote')).toBe(true); + expect(store.loadingConvs).toEqual(['local', 'remote']); + }); + + it('applyRemoteSnapshot removes remote convs missing from the snapshot', () => { + store.applyRemoteSnapshot(['a', 'b']); + store.applyRemoteSnapshot(['a']); + + expect(store.isRemote('a')).toBe(true); + expect(store.isRemote('b')).toBe(false); + expect(store.loadingConvs).toEqual(['a']); + }); + + it('applyRemoteSnapshot keeps local convs absent from the snapshot', () => { + store.markLocal('local'); + store.applyRemoteSnapshot(['remote']); + store.applyRemoteSnapshot([]); + + expect(store.isLocal('local')).toBe(true); + expect(store.loadingConvs).toEqual(['local']); + }); + + it('loadingConvs does not duplicate a conv that is both local and remote', () => { + store.markLocal('a'); + store.applyRemoteSnapshot(['a']); + + expect(store.loadingConvs).toEqual(['a']); + }); +}); diff --git a/tools/ui/tests/unit/mcp-override-fallback.test.ts b/tools/ui/tests/unit/mcp-override-fallback.test.ts index 47d6ac25363e..12ed6e4c4b44 100644 --- a/tools/ui/tests/unit/mcp-override-fallback.test.ts +++ b/tools/ui/tests/unit/mcp-override-fallback.test.ts @@ -46,7 +46,7 @@ describe('conversationsStore MCP override resolution', () => { // The settings store constructor bails in node env (no `browser`), // so seed the config directly. The shape mirrors what `loadConfig` // would build from localStorage. - const { settingsStore } = await import('$lib/stores/settings.svelte'); + const { settingsStore } = await import('$lib/stores/settings/index.svelte'); const raw = localStorage.getItem(CONFIG_LOCALSTORAGE_KEY) ?? '{}'; const saved = JSON.parse(raw) as Record; @@ -73,77 +73,77 @@ describe('conversationsStore MCP override resolution', () => { } it('inherits server.enabled when no conversation is active', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = null; - expect(conversationsStore.isMcpServerEnabledForChat('alpha')).toBe(false); - expect(conversationsStore.isMcpServerEnabledForChat('bravo')).toBe(true); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('alpha')).toBe(false); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('bravo')).toBe(true); }); it('inherits server.enabled on a newly created chat with no overrides', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = makeConversation(); // Empty override list: must fall back to global server.enabled, not all-off. - expect(conversationsStore.isMcpServerEnabledForChat('alpha')).toBe(false); - expect(conversationsStore.isMcpServerEnabledForChat('bravo')).toBe(true); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('alpha')).toBe(false); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('bravo')).toBe(true); }); it('inherits server.enabled on a newly created chat when overrides is undefined', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = makeConversation(undefined); - expect(conversationsStore.isMcpServerEnabledForChat('alpha')).toBe(false); - expect(conversationsStore.isMcpServerEnabledForChat('bravo')).toBe(true); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('alpha')).toBe(false); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('bravo')).toBe(true); }); it('uses explicit per-chat overrides, with defaults for non-overridden servers', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); // Override flips bravo off for this chat, alpha keeps its global default. conversationsStore.activeConversation = makeConversation([ { enabled: false, serverId: 'bravo' } ]); - expect(conversationsStore.isMcpServerEnabledForChat('alpha')).toBe(false); - expect(conversationsStore.isMcpServerEnabledForChat('bravo')).toBe(false); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('alpha')).toBe(false); + expect(conversationsStore.preferences.isMcpServerEnabledForChat('bravo')).toBe(false); }); it('getAllMcpServerOverrides returns a complete list merged from defaults', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = makeConversation([ { enabled: true, serverId: 'alpha' } ]); - expect(conversationsStore.getAllMcpServerOverrides()).toEqual([ + expect(conversationsStore.preferences.getAllMcpServerOverrides()).toEqual([ { enabled: true, serverId: 'alpha' }, { enabled: true, serverId: 'bravo' } ]); }); it('getAllMcpServerOverrides falls back to defaults when there are no explicit overrides', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = makeConversation(); - expect(conversationsStore.getAllMcpServerOverrides()).toEqual([ + expect(conversationsStore.preferences.getAllMcpServerOverrides()).toEqual([ { enabled: false, serverId: 'alpha' }, { enabled: true, serverId: 'bravo' } ]); }); it('getMcpServerOverride returns the global default when the server has no explicit override', async () => { - const { conversationsStore } = await import('$lib/stores/conversations.svelte'); + const { conversationsStore } = await import('$lib/stores/conversations/index.svelte'); conversationsStore.activeConversation = makeConversation([ { enabled: true, serverId: 'alpha' } ]); - expect(conversationsStore.getMcpServerOverride('bravo')).toEqual({ + expect(conversationsStore.preferences.getMcpServerOverride('bravo')).toEqual({ enabled: true, serverId: 'bravo' }); diff --git a/tools/ui/tests/unit/stream-resume.test.ts b/tools/ui/tests/unit/stream-resume.test.ts index 43d89272efd6..ce4eee9aa7ca 100644 --- a/tools/ui/tests/unit/stream-resume.test.ts +++ b/tools/ui/tests/unit/stream-resume.test.ts @@ -92,6 +92,67 @@ describe('ChatService stream resume', () => { expect(ChatService.getStreamState('conv-a')!.model).toBe('model-y'); }); + describe('throttled saves (per-chunk path)', () => { + // unique conversation ids: the throttle tracker is module state and + // outlives beforeEach's localStorage.clear() + let counter = 0; + + const freshConv = () => `conv-throttle-${++counter}`; + + it('writes immediately when no write was recorded for the conversation', () => { + const conv = freshConv(); + + ChatService.saveStreamStateThrottled(conv, 100); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(100); + }); + + it('holds a save pending when it lands inside the interval, flush forces it out', () => { + const conv = freshConv(); + + ChatService.saveStreamStateThrottled(conv, 100); + ChatService.saveStreamStateThrottled(conv, 200); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(100); + + ChatService.flushStreamState(conv); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(200); + }); + + it('flush is a no-op when nothing is pending', () => { + const conv = freshConv(); + + ChatService.saveStreamStateThrottled(conv, 100); + ChatService.flushStreamState(conv); + ChatService.flushStreamState(conv); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(100); + }); + + it('an immediate save resets the throttle window', () => { + const conv = freshConv(); + + ChatService.saveStreamStateThrottled(conv, 100); + ChatService.saveStreamState(conv, 150); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(150); + + ChatService.saveStreamStateThrottled(conv, 200); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(150); + + ChatService.flushStreamState(conv); + expect(ChatService.getStreamState(conv)!.bytesReceived).toBe(200); + }); + + it('clearStreamState drops the pending throttled state', () => { + const conv = freshConv(); + + ChatService.saveStreamStateThrottled(conv, 100); + ChatService.saveStreamStateThrottled(conv, 200); + ChatService.clearStreamState(conv); + expect(ChatService.getStreamState(conv)).toBeNull(); + + ChatService.flushStreamState(conv); + expect(ChatService.getStreamState(conv)).toBeNull(); + }); + }); + describe('resumeStreamIdentity', () => { it('appends the persisted model so the resume key matches the frozen POST identity', () => { ChatService.saveStreamState('conv-a', 10, 'model-x'); From 6b4fa88a6ce2429958ea4ee7c0223928e0979fa2 Mon Sep 17 00:00:00 2001 From: lhez Date: Thu, 20 Aug 2026 10:52:07 -0700 Subject: [PATCH 188/211] opencl: fix local size for norm (#27339) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index fbf7dadb906c..d169f33896b5 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -12830,7 +12830,10 @@ static void ggml_cl_norm(ggml_backend_t backend, const ggml_tensor * src0, const GGML_TENSOR_LOCALS(int, ne0, src0, ne); GGML_TENSOR_LOCALS(cl_ulong, nb0, src0, nb); - const int nth = MIN(64, ne00); + int nth = 1; + while (nth < ne00 && nth < 64) { + nth *= 2; + } cl_kernel kernel = backend_ctx->kernel_norm; From 6503355df0eb4f65875012523263c302fe0088c1 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Thu, 20 Aug 2026 10:58:35 -0700 Subject: [PATCH 189/211] opencl: fix q6_K flat mul_mat for Adreno A6x/A7x GPUs with older E031 compilers (#26476) * opencl: decline KV-convert flash_attn variants on Adreno A7X (compiler SIGSEGV) The Adreno 740 (A7X) compiler E031.41 crashes inside clBuildProgram when building the flash_attn programs whose KV path is mixed-type or dequantized: flash_attn_f32_f16, flash_attn_f32_q8_0, flash_attn_f32_q4_0. It is a driver crash rather than a compile-error return, so build_program_from_source_ex() cannot catch it. The uniform f32 and f16 programs build correctly. Decline the three KV-convert variants on the A7X in supports_op so they never lazy-compile; those attention layers run on the CPU backend instead. Same idiom as the existing Intel DK=512 and X1E carve-outs. test-backend-ops FLASH_ATTN_EXT on the 740: 226 OK / 0 FAIL, previously exit 139. Other parts are unaffected - the gate is dead code there. * opencl: fix q6_K flat mul_mat on older Adreno E031 compilers, gated kernel_mul_mv_q6_K_f32_flat produces ~10x-wrong output on the older Adreno E031 compilers while q4_K and q5_K are correct. Four codegen defects, each confirmed on-device against the CPU reference: 1. 64-bit ulong arithmetic is miscompiled, so every weight and scale read hit the wrong address - the primary cause, and why q5_K (int offsets) was unaffected. The block index is computed in int and widened only inside the pointer expression. 2. The vectorized dequant (int4/float4 bit-ops, convert_*4, dot()) is miscompiled; the 6-bit weights are reconstructed and the dot done scalar. 3. vload4 of the f32 activations is miscompiled; replaced by a scalar-indexed load. 4. The accumulation is miscompiled unless a side effect forces the partial sums to materialize. A printf under a guard the compiler cannot prove false acts as a zero-cost optimizer barrier; its placement is load-bearing. The defect tracks the compiler, not the GPU generation: it reproduces on E031.38 (Adreno 642L) and E031.41 (Adreno 740) and is fixed by E031.45 (Adreno 619), so the workarounds are gated on the compiler version. Where they are not needed they cost real throughput - 42.4 -> 35.1 GFLOPS on an Adreno 840 q6_K GEMV. The explicit compiler-type check is required, not redundant: newer_than_or_same() is false for every non-E031 compiler, so negating it alone would enable the workarounds on E17 and DX. test-backend-ops MUL_MAT is 919/919 on the Adreno 740, 642L, 619, 840 and 850; the 740 and 642L were 909/919 before. The 642L additionally needs the A6X per-kernel-program support to reach these tests at all. --- ggml/src/ggml-opencl/ggml-opencl.cpp | 42 +++++++- .../kernels/mul_mv_q6_k_f32_flat.cl | 96 ++++++++++++++++++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index d169f33896b5..49cd9fd35566 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -582,6 +582,8 @@ struct ggml_backend_opencl_context { bool adreno_use_bin_kernels; get_adreno_bin_kernel_func_t get_adreno_bin_kernel_func = nullptr; ggml_cl_compiler_version adreno_cl_compiler_version; + // The q6_K flat mul_mat codegen workarounds are needed by old E031 compilers only. + bool q6_k_flat_old_compiler; std::string kernel_compile_opts; // cached for lazy-compiled kernels. @@ -1931,8 +1933,14 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { #else const std::string kernel_src = read_file("mul_mv_q6_k_f32_flat.cl"); #endif + // The codegen workarounds in this kernel are a measured 13-20% loss on + // compilers that do not need them, so only the affected ones build them; + // everyone else gets the original source. + const std::string q6k_opts = backend_ctx->q6_k_flat_old_compiler + ? compile_opts + " -DADRENO_OLD_COMPILER=1" + : compile_opts; cl_program prog = - build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts); + build_program_from_source(backend_ctx, kernel_src.c_str(), q6k_opts); CL_CHECK((backend_ctx->kernel_mul_mv_q6_K_f32_flat = clCreateKernel(prog, "kernel_mul_mv_q6_K_f32_flat", &err), err)); CL_CHECK(clReleaseProgram(prog)); @@ -5917,6 +5925,16 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) { (backend_ctx->adreno_cl_compiler_version.type == E031 && backend_ctx->adreno_cl_compiler_version.major >= 47) || (backend_ctx->adreno_cl_compiler_version.type == DX && backend_ctx->adreno_cl_compiler_version.major >= 17); + // The q6_K flat mul_mat miscompile is a defect of the older E031 compilers, not a + // property of any GPU generation: it reproduces on E031.38 (Adreno 642L) and E031.41 + // (Adreno 740) and is fixed by E031.45 (Adreno 619). Gate on the compiler so parts + // that do not need the workarounds do not pay for them. The explicit type check is + // required: newer_than_or_same() is false for every non-E031 compiler, so negating it + // alone would enable the workarounds on E17/DX. + backend_ctx->q6_k_flat_old_compiler = + backend_ctx->adreno_cl_compiler_version.type == E031 && + !backend_ctx->adreno_cl_compiler_version.newer_than_or_same(E031, 45, 0, 0); + size_t ext_str_size; clGetDeviceInfo(device, CL_DEVICE_EXTENSIONS, 0, NULL, &ext_str_size); char *ext_buffer = (char *)alloca(ext_str_size + 1); @@ -7496,6 +7514,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F16; const bool is_f32_f16 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F32; + const bool is_f32_q8_0 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_Q8_0 && v->type == GGML_TYPE_Q8_0 && op->type == GGML_TYPE_F32 && dk % 32 == 0 && dv % 32 == 0; @@ -7503,6 +7522,21 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te v->type == GGML_TYPE_Q4_0 && op->type == GGML_TYPE_F32 && dk % 32 == 0 && dv % 32 == 0; + // A7X (Adreno 740, compiler E031.41) SIGSEGVs inside clBuildProgram + // building the flash_attn programs whose KV path is mixed-type or + // dequantized — f32_f16, q8_0, q4_0 (reproduced at DK=40 and DK=64; it + // is DK-independent). It is a driver crash, not codegen-wrong-output, so + // it cannot be caught in-process (fatal=false only handles clean compile + // errors). The uniform f16_f16 / f32_f32 programs compile fine on this + // compiler, so decline only the KV-convert variants; ggml then runs + // those (f16-KV / quant-KV) attention layers on the CPU backend. + // Negative compiler carve-out, same idiom as the Intel DK=512 decline + // below and the X1E driver-quirk guards. + if (backend_ctx && backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X && + (is_f32_f16 || is_f32_q8_0 || is_f32_q4_0)) { + return false; + } + // Asymmetric KV: host-dequants both sides to F32, uses f32 kernel. auto is_kv_type_ok = [](ggml_type t) { return t == GGML_TYPE_F16 || t == GGML_TYPE_F32 || @@ -20583,6 +20617,12 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int), &ne1)); CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &r2)); CL_CHECK(clSetKernelArg(kernel, 16, sizeof(int), &r3)); + // The optimizer-barrier arg exists only in the ADRENO_OLD_COMPILER build of + // this kernel; conformant compilers get the original 17-arg signature. + if (backend_ctx->q6_k_flat_old_compiler) { + cl_uchar q6k_mask = 0xFF; // never 0xFE in prod; see the kernel note + CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_uchar), &q6k_mask)); + } #else kernel = backend_ctx->kernel_mul_mv_q6_K_f32; diff --git a/ggml/src/ggml-opencl/kernels/mul_mv_q6_k_f32_flat.cl b/ggml/src/ggml-opencl/kernels/mul_mv_q6_k_f32_flat.cl index 57b90c05ae5f..2cca5335dd37 100644 --- a/ggml/src/ggml-opencl/kernels/mul_mv_q6_k_f32_flat.cl +++ b/ggml/src/ggml-opencl/kernels/mul_mv_q6_k_f32_flat.cl @@ -28,6 +28,13 @@ #define QK_K 256 +// ADRENO_OLD_COMPILER is defined by the host (-D) only for the Adreno E031 +// compilers older than E031.45, which miscompile several constructs this kernel +// used (confirmed on E031.38 and E031.41; E031.45 is clean). Every other +// compiler -- newer E031, E17, DX, Intel, and every non-Adreno device that +// builds this program -- takes the #else branches, which are the original +// source: the workarounds below cost ~13% on the q6_K flat n=1 GEMV where they +// are not needed. inline float block_q_6_K_dot_y_flat( global uchar * blk_ql, global uchar * blk_qh, @@ -37,6 +44,9 @@ inline float block_q_6_K_dot_y_flat( int ip, int is, int l0, +#if defined(ADRENO_OLD_COMPILER) + int dbg, +#endif float4 y0, float4 y1, float4 y2, @@ -48,10 +58,40 @@ inline float block_q_6_K_dot_y_flat( global uchar * q1 = blk_ql + ib*128 + q_offset_l; global uchar * q2 = q1 + QK_K/8; global uchar * qh = blk_qh + ib*64 + q_offset_h; - global char * sc = blk_scales + ib*16 + is; float dall = blk_d[ib]; +#if defined(ADRENO_OLD_COMPILER) + // The vectorized dequant (int4/float4 bit-ops, convert_*4, dot()) and vload4 + // are miscompiled here -> garbage weights. Reconstruct the 6-bit weights and + // take the dot product scalar. q4_K/q5_K flat already use scalar paths, which + // is why q6_K was the only flat GEMV that failed. + // Scales are SIGNED int8; read as uchar and sign-extend arithmetically so the + // result does not depend on whether the compiler treats `char` as signed. + global uchar * sc = (global uchar *)(blk_scales + ib*16 + is); + + int s0 = (int)sc[0] - 256*(sc[0] >> 7); + int s2 = (int)sc[2] - 256*(sc[2] >> 7); + int s4 = (int)sc[4] - 256*(sc[4] >> 7); + int s6 = (int)sc[6] - 256*(sc[6] >> 7); + + // one 6-bit weight: low/high nibble of a ql byte OR'd with a 2-bit qh plane + // (plane p in {0,1,2,3} selects qh bits 2p..2p+1) placed at bits 4-5, minus 32. + #define Q6W(qb, sh, hb, p) ((float)((((int)(qb) >> (sh)) & 15) | ((((int)(hb) >> (2*(p))) & 3) << 4)) - 32.f) + + float d0 = y0.s0*Q6W(q1[0],0,qh[0],0) + y0.s1*Q6W(q1[1],0,qh[1],0) + y0.s2*Q6W(q1[2],0,qh[2],0) + y0.s3*Q6W(q1[3],0,qh[3],0); + float d1 = y1.s0*Q6W(q2[0],0,qh[0],1) + y1.s1*Q6W(q2[1],0,qh[1],1) + y1.s2*Q6W(q2[2],0,qh[2],1) + y1.s3*Q6W(q2[3],0,qh[3],1); + float d2 = y2.s0*Q6W(q1[0],4,qh[0],2) + y2.s1*Q6W(q1[1],4,qh[1],2) + y2.s2*Q6W(q1[2],4,qh[2],2) + y2.s3*Q6W(q1[3],4,qh[3],2); + float d3 = y3.s0*Q6W(q2[0],4,qh[0],3) + y3.s1*Q6W(q2[1],4,qh[1],3) + y3.s2*Q6W(q2[2],4,qh[2],3) + y3.s3*Q6W(q2[3],4,qh[3],3); + #undef Q6W + + if (dbg) printf("HELPER dall=%f s=[%d %d %d %d] d=[%f %f %f %f] ql0=%d qh0=%d y00=%f\n", + dall, s0, s2, s4, s6, d0, d1, d2, d3, (int)q1[0], (int)qh[0], y0.s0); + + return dall * (d0 * s0 + d1 * s2 + d2 * s4 + d3 * s6); +#else + global char * sc = blk_scales + ib*16 + is; + // Vectorized loads: 3 uchar4 weight loads instead of 12 scalar byte reads. // q_offset_l/h are 4-aligned, so these are aligned vector loads. uchar4 q1v = vload4(0, q1); @@ -72,6 +112,7 @@ inline float block_q_6_K_dot_y_flat( return dall * (dot(y0, w0) * sc[0] + dot(y1, w1) * sc[2] + dot(y2, w2) * sc[4] + dot(y3, w3) * sc[6]); +#endif } #undef N_DST @@ -113,6 +154,11 @@ kernel void kernel_mul_mv_q6_K_f32_flat( int ne1, int r2, int r3 +#if defined(ADRENO_OLD_COMPILER) + , + uchar q6k_mask // runtime 0xFF; the host passes it so the compiler cannot + // constant-fold the printf guards below into nothing +#endif ) { src1 = (global float*)((global char*)src1 + offset1); dst = (global float*)((global char*)dst + offsetd); @@ -128,6 +174,22 @@ kernel void kernel_mul_mv_q6_K_f32_flat( int first_row = (N_SIMDGROUP * r0 + get_sub_group_id()) * N_DST; +#if defined(ADRENO_OLD_COMPILER) + // 64-bit `ulong` integer arithmetic is miscompiled here -> the base-pointer byte + // offsets came out wrong, so EVERY weight/scale read hit the wrong address. This + // was the primary cause of the q6_K flat failure (q5_K uses int offsets and is + // unaffected). Compute the block index in `int` and widen to `ulong` only inside + // the pointer expression: the byte offset stays 64-bit, but there is no ulong + // arithmetic chain to miscompile. The int index would overflow past ~2^31 blocks, + // which no realistic weight reaches -- but that is a narrowing, so keep it off the + // conformant path, which retains full ulong arithmetic. + int offset_src0 = first_row*nb + (i12/r2)*(nb*ne01) + (i13/r3)*(nb*ne01*ne02); + + global uchar * blk_ql = (global uchar *) src0_ql + (ulong)offset_src0 * 128; + global uchar * blk_qh = (global uchar *) src0_qh + (ulong)offset_src0 * 64; + global char * blk_scales = (global char *) src0_s + (ulong)offset_src0 * 16; + global half * blk_d = (global half *) src0_d + offset_src0; +#else ulong offset_src0 = first_row*nb + (i12/r2)*(nb*ne01) + (i13/r3)*(nb*ne01*ne02); ulong offset_src0_ql = offset_src0 * 128; ulong offset_src0_qh = offset_src0 * 64; @@ -138,6 +200,7 @@ kernel void kernel_mul_mv_q6_K_f32_flat( global uchar * blk_qh = (global uchar *) src0_qh + offset_src0_qh; global char * blk_scales = (global char *) src0_s + offset_src0_s; global half * blk_d = (global half *) src0_d + offset_src0_d; +#endif global float * yy = (global float *) src1 + r1*ne10 + im*ne00*ne1; int tid = get_sub_group_local_id()%(N_SIMDWIDTH/BLOCK_STRIDE); // within-super-block part, 0..15 @@ -155,24 +218,55 @@ kernel void kernel_mul_mv_q6_K_f32_flat( for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) { global float * y = yy + ib * QK_K + 128*ip + l0; +#if defined(ADRENO_OLD_COMPILER) + // vload4 of f32 is miscompiled here; index the lanes scalar instead. + float4 y0 = (float4)(y[ 0], y[ 1], y[ 2], y[ 3]); + float4 y1 = (float4)(y[32], y[33], y[34], y[35]); + float4 y2 = (float4)(y[64], y[65], y[66], y[67]); + float4 y3 = (float4)(y[96], y[97], y[98], y[99]); +#else float4 y0 = vload4(0, y + 0); float4 y1 = vload4(0, y + 32); float4 y2 = vload4(0, y + 64); float4 y3 = vload4(0, y + 96); +#endif for (int row = 0; row < N_DST; row++) { if (first_row + row < ne01) { +#if defined(ADRENO_OLD_COMPILER) + int dbg = (q6k_mask==0xFE && r0==0 && r1==0 && im==0 && row==0 && ib==0 && + ne00==256 && ne01==16 && get_sub_group_local_id()==0) ? 1 : 0; + sumf[row] += block_q_6_K_dot_y_flat( + blk_ql + row*nb*128, blk_qh + row*nb*64, blk_scales + row*nb*16, blk_d + row*nb, + ib, ip, is, l0, dbg, y0, y1, y2, y3); +#else sumf[row] += block_q_6_K_dot_y_flat( blk_ql + row*nb*128, blk_qh + row*nb*64, blk_scales + row*nb*16, blk_d + row*nb, ib, ip, is, l0, y0, y1, y2, y3); +#endif } } } +#if defined(ADRENO_OLD_COMPILER) + // Optimizer barrier. This compiler drops the sumf partials unless a side effect + // forces them to materialize. q6k_mask is a kernel arg the compiler cannot prove + // is never 0xFE (the host always passes 0xFF), so the printf survives compilation + // but never executes. FRAGILE: the exact set and placement of these guarded + // printfs is load-bearing on E031.41 -- removing any one re-breaks q6_K. + if (q6k_mask==0xFE && r0==0 && r1==0 && im==0 && ne00==256 && ne01==16 && get_sub_group_local_id()<16) { + printf("Q6KLANE lane=%d ip=%d il=%d is=%d l0=%d sumf0=%f\n", + get_sub_group_local_id(), ip, il, is, l0, sumf[0]); + } +#endif for (int row = 0; row < N_DST; row++) { float tot = sub_group_reduce_add(sumf[row]); if (get_sub_group_local_id() == 0 && first_row + row < ne01) { dst[r1*ne0 + im*ne0*ne1 + first_row + row] = tot; +#if defined(ADRENO_OLD_COMPILER) + if (q6k_mask==0xFE && r0==0 && r1==0 && im==0 && row==0 && ne00==256 && ne01==16) + printf("Q6KTOT tot=%f\n", tot); +#endif } } } From a30273376ef669023334fc20ad02ae4ed8196a65 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 20 Aug 2026 21:31:29 +0300 Subject: [PATCH 190/211] metal : clamp K extent in tensor API mat-mat kernel for K not a multiple of 32 (#27450) The Tensor API mat-mat path of kernel_mul_mm (GGML_METAL_HAS_TENSOR) fed a static K=32 tile to the matmul2d op on every iteration. On the last, partial K tile (ne00 % 32 != 0) the src1 slice extends past the K extent of the tensor, and the op reads those out-of-bounds elements (undefined behavior per the MSL specification, section 2.22.2). Depending on stale memory contents, this corrupted the result or produced NaN. Make the matmul2d op use dynamic_extent for K, and clamp the K extent of both operand tensor views to the remaining valid K range (min(32, K - loop_k)) per iteration, so the op reads exactly the valid K range on every iteration (mirroring the tail handling of the MPP matmul2d examples). On K-aligned inputs the clamp degenerates to the full 32-wide tile: the only difference from the static-K op is that the dynamic-K op derives K from the operand extents and edge-checks the tile against the tensor extents (a handful of integer ops per iteration). Add test-backend-ops MUL_MAT cases with K not a multiple of 32 to exercise the unaligned K path. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal.metal | 15 +++++++++++---- tests/test-backend-ops.cpp | 8 ++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 949931c8dc5a..27f97b5e0798 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -10365,9 +10365,12 @@ kernel void kernel_mul_mm( auto tB = tensor(ptrB, dextents(K, N), array({1, strideB})); // Configure matmul operation + // note: K is dynamic_extent (clamped to the valid range in PHASE 2), since a static + // N_MM_NK_TOTAL K tile would read src1 out of bounds when K % N_MM_NK_TOTAL != 0 + // ref: https://github.com/ggml-org/llama.cpp/pull/27064 mpp::tensor_ops::matmul2d< mpp::tensor_ops::matmul2d_descriptor( - NRB, NRA, N_MM_NK_TOTAL, false, true, true, + NRB, NRA, static_cast(dynamic_extent), false, true, true, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate), execution_simdgroups> mm; @@ -10419,10 +10422,14 @@ kernel void kernel_mul_mm( threadgroup_barrier(mem_flags::mem_threadgroup); // === PHASE 2: Tensor matmul === - auto mA = tA.slice(0, 0); - auto mB = tB.slice(loop_k, rb); + // Clamp the K extent of both operand tensors to the remaining valid K range so + // the dynamic-K op never reads past the K extent of src1 (or the staged A tile). + const int kExt = min(N_MM_NK_TOTAL, K - loop_k); - mm.run(mB, mA, cT); + auto tAv = tensor(sa, dextents(kExt, NRA), array({1, N_MM_NK_TOTAL})); + auto tBv = tensor(ptrB + loop_k + rb * strideB, dextents(kExt, N - rb), array({1, strideB})); + + mm.run(tBv, tAv, cT); threadgroup_barrier(mem_flags::mem_threadgroup); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 17098825bc81..89b954a7d1e0 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9298,6 +9298,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6, 4096, 5120, {1, 1}, {1, 1})); + // K not a multiple of 32 + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 65, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 64, 32, 80, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 32, 80, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 588, {1, 1}, {1, 1})); // 14*14*3, e.g. conv_2d im2col + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 64, 32, 80, {4, 1}, {1, 1})); + #if 0 // test the mat-mat path for Metal for (int k = 1; k < 512; ++k) { From 0e1d9185c5fe82e905d1f5ae6b2e5dcd607a8dfd Mon Sep 17 00:00:00 2001 From: Eve <139727413+netrunnereve@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:01:32 +0000 Subject: [PATCH 191/211] ci: use shell script to check cmake pkg (#27414) * use regular script to build cmake pkg * use old grep without perl --- .github/workflows/build-cmake-pkg.yml | 42 ++++++++++++--------------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/.github/workflows/build-cmake-pkg.yml b/.github/workflows/build-cmake-pkg.yml index 0e4069ce3526..c44fba2c6953 100644 --- a/.github/workflows/build-cmake-pkg.yml +++ b/.github/workflows/build-cmake-pkg.yml @@ -27,30 +27,26 @@ jobs: cmake --install build --prefix "$PREFIX" --config Release export LLAMA_CONFIG="$PREFIX"/lib/cmake/llama/llama-config.cmake - tclsh <<'EOF' - set build(commit) [string trim [exec git rev-parse --short HEAD]] - set build(number) [string trim [exec git rev-list --count HEAD]] - - set cmakelists [read [open "CMakeLists.txt" r]] - regexp {set\(LLAMA_VERSION_MAJOR\s+(\d+)\)} $cmakelists -> major - regexp {set\(LLAMA_VERSION_MINOR\s+(\d+)\)} $cmakelists -> minor - regexp {set\(LLAMA_VERSION_PATCH\s+(\d+)\)} $cmakelists -> patch - set build(version) "$major.$minor.$patch" - - set llamaconfig [read [open "$env(LLAMA_CONFIG)" r]] - set checks [list "set\\(LLAMA_VERSION \\s+$build(version)\\)" \ - "set\\(LLAMA_BUILD_COMMIT\\s+$build(commit)\\)" \ - "set\\(LLAMA_BUILD_NUMBER\\s+$build(number)\\)"] - - puts -nonewline "Checking llama-config.cmake version... " - foreach check $checks { - if {![regexp -expanded -- $check $llamaconfig]} { - puts "\"$check\" failed!" + build_commit=$(git rev-parse --short HEAD | xargs) + build_number=$(git rev-list --count HEAD | xargs) + + major=$(grep -oE "set\(LLAMA_VERSION_MAJOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$") + minor=$(grep -oE "set\(LLAMA_VERSION_MINOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$") + patch=$(grep -oE "set\(LLAMA_VERSION_PATCH[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$") + build_version="$major.$minor.$patch" + + checks=("set\(LLAMA_VERSION[[:space:]]+$build_version\)" + "set\(LLAMA_BUILD_COMMIT[[:space:]]+$build_commit\)" + "set\(LLAMA_BUILD_NUMBER[[:space:]]+$build_number\)") + + for check in "${checks[@]}"; do + if ! grep -qE "$check" "$LLAMA_CONFIG"; then + echo "Checking llama-config.cmake version... \"$check\" failed!" exit 1 - } - } - puts "success." - EOF + fi + done + + echo "Checking llama-config.cmake version... success." cd examples/simple-cmake-pkg cmake -S . -B build -DCMAKE_PREFIX_PATH="$PREFIX"/lib/cmake From 749f688fcaa4c472ec034b08cb8a907c45cfaa02 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 21 Aug 2026 00:36:57 +0200 Subject: [PATCH 192/211] ggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon (#27345) * ggml: support ggml_rope_set_offset on opencl, sycl, wgpu, hexagon * rm inplace optimization --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 5 +- ggml/src/ggml-hexagon/htp/rope-ops.c | 22 ++-- ggml/src/ggml-opencl/ggml-opencl.cpp | 11 +- ggml/src/ggml-opencl/kernels/rope.cl | 92 +++++++++-------- ggml/src/ggml-sycl/ggml-sycl.cpp | 2 - ggml/src/ggml-sycl/rope.cpp | 106 +++++++++++--------- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 8 +- ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl | 18 ++-- 8 files changed, 152 insertions(+), 112 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index b262a73d950e..e8a5009b381b 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3180,8 +3180,9 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const int32_t * op_params = &op->op_params[0]; - if (op_params[15] != 0) { - return false; // FIXME: support ggml_rope_set_offset + // ggml_rope_set_offset: HVX kernels need a VLEN-aligned window start (32 f32 elems) + if (op_params[15] % 32 != 0) { + return false; } int mode = op_params[2]; diff --git a/ggml/src/ggml-hexagon/htp/rope-ops.c b/ggml/src/ggml-hexagon/htp/rope-ops.c index 5bc7d74f5e21..6c689824934f 100644 --- a/ggml/src/ggml-hexagon/htp/rope-ops.c +++ b/ggml/src/ggml-hexagon/htp/rope-ops.c @@ -53,6 +53,7 @@ struct htp_rope_context { int32_t n_dims; + int32_t n_offs; int32_t mode; int32_t n_ctx_orig; int32_t sections[4]; @@ -405,32 +406,40 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict static void inline rope_basic_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src, uint32_t nr, uint32_t ne0, const float * restrict theta_cache) { + const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op) #pragma unroll(4) for (uint32_t i = 0; i < nr; i++) { float * d = (float *) (dst + i * rctx->dst_row_size_aligned); float * s = (float *) (src + i * rctx->src0_row_size_aligned); - hvx_rope_f32_aa(d, s, rctx->n_dims, theta_cache); + hvx_rope_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache); // fill the remain channels with data from src tensor - if (rctx->n_dims < ne0) { - hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims); + if (n_offs > 0) { + hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs); + } + if (n_offs + rctx->n_dims < ne0) { + hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims); } } } static void inline rope_neox_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src, uint32_t nr, uint32_t ne0, const float * restrict theta_cache) { + const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op) #pragma unroll(4) for (uint32_t i = 0; i < nr; i++) { float * d = (float *) (dst + i * rctx->dst_row_size_aligned); float * s = (float *) (src + i * rctx->src0_row_size_aligned); - hvx_rope_neox_f32_aa(d, s, rctx->n_dims, theta_cache); + hvx_rope_neox_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache); // fill the remain channels with data from src tensor - if (rctx->n_dims < ne0) { - hvx_copy_f32_uu((uint8_t *)(d + rctx->n_dims), (uint8_t *)(s + rctx->n_dims), ne0 - rctx->n_dims); + if (n_offs > 0) { + hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs); + } + if (n_offs + rctx->n_dims < ne0) { + hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims); } } } @@ -673,6 +682,7 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) { rctx.n_dims = ((const int32_t *) op_params)[1]; rctx.mode = ((const int32_t *) op_params)[2]; rctx.n_ctx_orig = ((const int32_t *) op_params)[4]; + rctx.n_offs = ((const int32_t *) op_params)[15]; memcpy(&rctx.freq_base, (int32_t *) op_params + 5, sizeof(float)); memcpy(&rctx.freq_scale, (int32_t *) op_params + 6, sizeof(float)); diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 49cd9fd35566..26f952a173ad 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7434,9 +7434,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_DIAG_MASK_INF: return op->ne[3] == 1; case GGML_OP_ROPE: { - if (((const int32_t *) op->op_params)[15] != 0) { - return false; // FIXME: support ggml_rope_set_offset - } const int mode = ((const int32_t *) op->op_params)[2]; const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; const bool is_vision = mode == GGML_ROPE_TYPE_VISION; @@ -23910,6 +23907,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const const int n_dims = ((int *) dst->op_params)[1]; const int mode = ((int *) dst->op_params)[2]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; float freq_base; float freq_scale; @@ -23938,6 +23936,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const if (is_vision) { GGML_ASSERT(n_dims == ne00/2); + GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row } cl_kernel kernel; @@ -24029,6 +24028,12 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const if (is_mrope && !is_vision) { CL_CHECK(clSetKernelArg(kernel, 34, sizeof(int), &is_imrope)); } + // norm and neox have n_offs after beta_slow, mrope has it after is_imrope + if (!is_mrope && !is_vision) { + CL_CHECK(clSetKernelArg(kernel, 33, sizeof(int), &n_offs)); + } else if (is_mrope && !is_vision) { + CL_CHECK(clSetKernelArg(kernel, 35, sizeof(int), &n_offs)); + } size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03}; size_t local_work_size[] = {(size_t)nth, 1, 1}; diff --git a/ggml/src/ggml-opencl/kernels/rope.cl b/ggml/src/ggml-opencl/kernels/rope.cl index 82f4cd87407d..27fdbbbc4ff1 100644 --- a/ggml/src/ggml-opencl/kernels/rope.cl +++ b/ggml/src/ggml-opencl/kernels/rope.cl @@ -75,7 +75,8 @@ kernel void kernel_rope_norm_f32( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -94,14 +95,15 @@ kernel void kernel_rope_norm_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - float theta = theta_base * pow(freq_base, inv_ndims*i0); + float theta = theta_base * pow(freq_base, inv_ndims*iw); float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); @@ -154,7 +156,8 @@ kernel void kernel_rope_norm_f16( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -173,14 +176,15 @@ kernel void kernel_rope_norm_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - float theta = theta_base * pow(freq_base, inv_ndims*i0); + float theta = theta_base * pow(freq_base, inv_ndims*iw); float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); @@ -233,7 +237,8 @@ kernel void kernel_rope_neox_f32( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -252,17 +257,18 @@ kernel void kernel_rope_neox_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -312,7 +318,8 @@ kernel void kernel_rope_neox_f16( float ext_factor, float attn_factor, float beta_fast, - float beta_slow + float beta_slow, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -331,17 +338,18 @@ kernel void kernel_rope_neox_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -393,7 +401,8 @@ kernel void kernel_rope_multi_f32( float beta_fast, float beta_slow, int4 sections, - int is_imrope + int is_imrope, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -414,10 +423,11 @@ kernel void kernel_rope_multi_f32( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const int sector = (i0 / 2) % sect_dims; + const int sector = ic % sect_dims; float theta_base = 0.0f; if (is_imrope) { @@ -445,14 +455,14 @@ kernel void kernel_rope_multi_f32( } } - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global float * src = (global float *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global float * dst_data = (global float *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; @@ -504,7 +514,8 @@ kernel void kernel_rope_multi_f16( float beta_fast, float beta_slow, int4 sections, - int is_imrope + int is_imrope, + int n_offs ) { src0 = (global void*)((global char*)src0 + offset0); src1 = (global int*)((global char*)src1 + offset1); @@ -525,10 +536,11 @@ kernel void kernel_rope_multi_f16( float inv_ndims = -1.f/n_dims; for (int i0 = 2*get_local_id(0); i0 < ne0; i0 += 2*get_local_size(0)) { - if (i0 < n_dims) { - int ic = i0/2; + if (i0 >= n_offs && i0 < n_offs + n_dims) { + int iw = i0 - n_offs; // relative idx + int ic = iw/2; - const int sector = (i0 / 2) % sect_dims; + const int sector = ic % sect_dims; float theta_base = 0.0f; if (is_imrope) { @@ -556,14 +568,14 @@ kernel void kernel_rope_multi_f16( } } - const float theta = theta_base * pow(freq_base, inv_ndims*i0); + const float theta = theta_base * pow(freq_base, inv_ndims*iw); const float freq_factor = src2 != src0 ? src2[ic] : 1.0f; - float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor); + float2 cos_sin_theta = rope_yarn(theta/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor); - global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + ic*nb00); - global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + ic*nb0); + global half * src = (global half *)((global char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + (n_offs + ic)*nb00); + global half * dst_data = (global half *)((global char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + (n_offs + ic)*nb0); const float x0 = src[0]; const float x1 = src[n_dims/2]; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index c7434a6bdbac..57aae9011d1f 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6242,8 +6242,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons } case GGML_OP_ROPE: case GGML_OP_ROPE_BACK: - // FIXME: support ggml_rope_set_offset - return ((const int32_t *) op->op_params)[15] == 0; case GGML_OP_IM2COL: case GGML_OP_IM2COL_3D: case GGML_OP_UPSCALE: diff --git a/ggml/src/ggml-sycl/rope.cpp b/ggml/src/ggml-sycl/rope.cpp index 9d83a1e9fa09..b6d22559d18c 100644 --- a/ggml/src/ggml-sycl/rope.cpp +++ b/ggml/src/ggml-sycl/rope.cpp @@ -41,7 +41,7 @@ template static void rope_norm(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -78,19 +78,21 @@ static void rope_norm(const T *x, D *dst, const int ne00, const int ne01, ggml_sycl_memcpy_1<4>(dst + idst, &v); } }; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { store_coaelsced(x[ix + 0], x[ix + 1]); return; } - const float theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f; + const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); + + const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, i0, + rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); const float x0 = x[ix + 0]; @@ -104,7 +106,7 @@ template static void rope_neox(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -132,35 +134,38 @@ static void rope_neox(const T *x, D *dst, const int ne00, const int ne01, idst += row_indices[i2] * set_rows_stride; } - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { dst[idst + i0 / 2 + 0] = ggml_sycl_cast(x[ix + i0 / 2 + 0]); dst[idst + i0 / 2 + 1] = ggml_sycl_cast(x[ix + i0 / 2 + 1]); return; } - const float theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f; + const float theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); + + const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, i0, + rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims / 2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs / 2 + 0]; + const float x1 = x[ix + n_offs / 2 + n_dims / 2]; - dst[idst + 0] = ggml_sycl_cast(x0 * cos_theta - x1 * sin_theta); - dst[idst + n_dims / 2] = ggml_sycl_cast(x0 * sin_theta + x1 * cos_theta); + dst[idst + n_offs / 2 + 0] = ggml_sycl_cast(x0 * cos_theta - x1 * sin_theta); + dst[idst + n_offs / 2 + n_dims / 2] = ggml_sycl_cast(x0 * sin_theta + x1 * cos_theta); } template static void rope_multi(const T *x, T *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, - const int n_dims, const int32_t *pos, + const int n_dims, const int n_offs, const int32_t *pos, const float freq_scale, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float theta_scale, const float *freq_factors, @@ -183,54 +188,57 @@ static void rope_multi(const T *x, T *dst, const int ne00, const int ne01, int idst = i0 / 2 + i1 * s1 + i2 * s2 + i3 * s3; const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { dst[idst + i0 / 2 + 0] = x[ix + i0 / 2 + 0]; dst[idst + i0 / 2 + 1] = x[ix + i0 / 2 + 1]; return; } + const int iw = i0 - n_offs; // relative idx + const int sect_dims = sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3]; const int sec_w = sections.v[1] + sections.v[0]; - const int sector = (i0 / 2) % sect_dims; + const int sector = (iw / 2) % sect_dims; float theta_base = 0.0; if (is_imrope) { if (sector % 3 == 1 && sector < 3 * sections.v[1]) { // h - theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w - theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t - theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); } else { - theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f); } } else { if (sector < sections.v[0]) { - theta_base = pos[i2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sections.v[0] && sector < sec_w) { - theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sec_w && sector < sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * dpct::pow(theta_scale, iw / 2.0f); } else if (sector >= sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * dpct::pow(theta_scale, iw / 2.0f); } } - const float freq_factor = has_ff ? freq_factors[i0 / 2] : 1.0f; + const float freq_factor = has_ff ? freq_factors[iw / 2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, i0, + rope_yarn(theta_base / freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims / 2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs / 2 + 0]; + const float x1 = x[ix + n_offs / 2 + n_dims / 2]; - dst[idst + 0] = x0 * cos_theta - x1 * sin_theta; - dst[idst + n_dims / 2] = x0 * sin_theta + x1 * cos_theta; + dst[idst + n_offs / 2 + 0] = x0 * cos_theta - x1 * sin_theta; + dst[idst + n_offs / 2 + n_dims / 2] = x0 * sin_theta + x1 * cos_theta; } template @@ -293,7 +301,7 @@ static void rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const int64_t *row_indices, @@ -313,7 +321,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_norm( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } else { @@ -323,7 +331,7 @@ rope_norm_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_norm( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } @@ -334,7 +342,7 @@ static void rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const int64_t *row_indices, @@ -354,7 +362,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_neox( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } else { @@ -364,7 +372,7 @@ rope_neox_sycl(const T *x, D *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_neox( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); }); } @@ -375,7 +383,7 @@ static void rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, const int ne02, const int s01, const int s02, const int s03, const int s1, const int s2, const int s3, const int n_dims, - const int nr, const int32_t *pos, const float freq_scale, + const int n_offs, const int nr, const int32_t *pos, const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor, const rope_corr_dims corr_dims, const float *freq_factors, const mrope_sections sections, @@ -395,7 +403,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_multi( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); }); } else { @@ -405,7 +413,7 @@ rope_multi_sycl(const T *x, T *dst, const int ne00, const int ne01, GGML_UNUSED(item_ct1); rope_multi( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, - pos, freq_scale, ext_factor, attn_factor, corr_dims, + n_offs, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); }); } @@ -497,6 +505,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, const int n_dims = ((int32_t *)dst->op_params)[1]; const int mode = ((int32_t *)dst->op_params)[2]; const int n_ctx_orig = ((int32_t *)dst->op_params)[4]; + const int n_offs = ((int32_t *)dst->op_params)[15]; mrope_sections sections; float freq_base; @@ -526,6 +535,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (is_vision) { GGML_ASSERT(n_dims == ne00 / 2); + GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row } const int32_t *pos = (const int32_t *)src1_d; @@ -545,19 +555,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_neox_sycl( (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, - s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_neox_sycl( (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02, - s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_neox_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else { @@ -568,13 +578,13 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32) { rope_multi_sycl((const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, s02, s03, s1, s2, - s3, n_dims, nr, pos, freq_scale, freq_base, + s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope, stream); } else if (src0->type == GGML_TYPE_F16) { rope_multi_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope, stream); } else { @@ -602,19 +612,19 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_norm_sycl( (const float *)src0_d, (float *)dst_d, ne00, ne01, ne02, s01, - s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_norm_sycl( (const float *)src0_d, (sycl::half *)dst_d, ne00, ne01, ne02, - s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_norm_sycl( (const sycl::half *)src0_d, (sycl::half *)dst_d, ne00, ne01, - ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, + ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, set_rows_stride, stream); } else { diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 4367f9a6109e..2434848a55a8 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -2714,6 +2714,7 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx, const int n_dims = ((int32_t *) dst->op_params)[1]; const int mode = ((int32_t *) dst->op_params)[2]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; float freq_base; float freq_scale; @@ -2762,7 +2763,8 @@ static webgpu_encoded_op ggml_webgpu_rope(webgpu_context & ctx, (uint32_t) sections[0], (uint32_t) sections[1], (uint32_t) sections[2], - (uint32_t) sections[3] + (uint32_t) sections[3], + (uint32_t) n_offs }; std::vector entries = { ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0), @@ -4472,9 +4474,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const supports_op = (op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32) && ggml_is_contiguous_rows(src0); break; case GGML_OP_ROPE: - // FIXME: support ggml_rope_set_offset - supports_op = - (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ((const int32_t *) op->op_params)[15] == 0; + supports_op = op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16; break; case GGML_OP_GLU: switch (ggml_get_glu_op(op)) { diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl index 1c874e14240e..6ff53088c461 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/rope.wgsl @@ -38,7 +38,8 @@ struct Params { sections0: u32, sections1: u32, sections2: u32, - sections3: u32 + sections3: u32, + n_offs: u32 }; @group(0) @binding(0) @@ -126,7 +127,8 @@ fn rope_yarn(theta_extrap: f32, i: u32) -> vec2 { fn pair_base(i0: u32, div_2: bool) -> u32 { if (div_2) { - return i0 / 2; + // first channel of the rotated pair: n_offs + (i0 - n_offs)/2 + return i0 / 2 + params.n_offs / 2; } else { return i0; } @@ -165,20 +167,22 @@ fn main(@builtin(global_invocation_id) gid: vec3) { let i_src_row = params.offset_src0 + i3 * params.stride_src03 + i2 * params.stride_src02 + i1 * params.stride_src01; let i_dst_row = params.offset_dst + i3 * params.stride_dst3 + i2 * params.stride_dst2 + i1 * params.stride_dst1; - if (i0 >= params.n_dims && !is_vision) { + if ((i0 < params.n_offs || i0 >= params.n_offs + params.n_dims) && !is_vision) { let i_src = i_src_row + i0; let i_dst = i_dst_row + i0; rotate(i_dst, i_dst + 1, f32(src0[i_src]), f32(src0[i_src + 1])); return; } + let iw = i0 - params.n_offs; // relative idx + var theta_base_mult: u32 = 0; - var theta_scale_pwr: u32 = i0 / 2; + var theta_scale_pwr: u32 = iw / 2; if (is_mrope) { let sect_dims = params.sections0 + params.sections1 + params.sections2 + params.sections3; let sec_w = params.sections1 + params.sections0; let sec_e = params.sections2 + sec_w; - let sector = (i0 / 2) % sect_dims; + let sector = (iw / 2) % sect_dims; if (is_imrope) { if (sector % 3 == 1 && sector < 3 * params.sections1) { theta_base_mult = 1; @@ -203,7 +207,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } else if (sector >= sec_e) { if (is_vision) { theta_scale_pwr = sector - sec_e; - theta_scale_pwr = (i0 / 2) % sec_e; + theta_scale_pwr = (iw / 2) % sec_e; } theta_base_mult = 3; } else if (is_vision) { @@ -212,7 +216,7 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } } let theta_base = f32(src1[params.offset_src1 + i2 + params.ne2 * theta_base_mult]) * pow(params.theta_scale, f32(theta_scale_pwr)); - let thetas = rope_yarn(theta_base/freq_factor(i0), i0); + let thetas = rope_yarn(theta_base/freq_factor(iw), iw); let i_src = i_src_row + pair_base(i0, is_neox || is_mrope || is_vision); let i_dst = i_dst_row + pair_base(i0, is_neox || is_mrope || is_vision); From a298422da78eb75e440a7de0ca408af64d323d93 Mon Sep 17 00:00:00 2001 From: vk <89937361+itsvedantkumar@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:06:59 +0530 Subject: [PATCH 193/211] docs: fix typos in ET.md (#27457) --- docs/backend/ET.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/backend/ET.md b/docs/backend/ET.md index 8d9ba12c822d..8ebc15fb7d43 100644 --- a/docs/backend/ET.md +++ b/docs/backend/ET.md @@ -116,7 +116,7 @@ in inline assembler. Most kernels are very naive with lots of low hanging fruits left: > [!IMPORTANT] -> Several assembly instructions emmited by the compiler are not implemented +> Several assembly instructions emitted by the compiler are not implemented > in hardware and software emulation in firmware is not ready yet. > Eventually firmware will transparently trap unimplemented instructions > and will emulate them inside exception handler. Until then, kernel @@ -138,12 +138,12 @@ Most kernels are very naive with lots of low hanging fruits left: > kernel build process. Feel free to take ideas/code from there or try linking > it in. -Before commiting any changes to operations and/or kernels, don't forget +Before committing any changes to operations and/or kernels, don't forget to update supported ops reports (instructions at `docs/ops.md`). When logging is enabled (e.g. by setting `--log-file` cli param), each compute kernel run outputs a line with -pipe-delimited key-value pairs containing kernel level performance infomation. +pipe-delimited key-value pairs containing kernel level performance information. Line is prefixed with `ET_PERF`: ``` @@ -160,7 +160,7 @@ to `GGML_ET_PROFILE/et_runtime_trace.json` and `GGML_ET_PROFILE/kernel_map` on e ### Uberkernel -The in-knernel implementaiton of device dispatch/kernel fusion. The ET SDK has a non-trivial op-to-op gap. `Uberkernel` (name taken from the original Esperanto AI's compiler) +The in-kernel implementation of device dispatch/kernel fusion. The ET SDK has a non-trivial op-to-op gap. `Uberkernel` (name taken from the original Esperanto AI's compiler) dispatches multiple already existing kernel implementations with device side synchronization. Due to the processor's design, there is no natural memory visibility horizon between sub-kernel invocations. This makes uberkernel much more difficult to develop and debug. Currently Uberkerel is hidden begind the `GGML_ET_UBERKERNEL` environment variable and is disabled by default. Setting it to 1 enables it and provides significant performance improvements but is only From b2e5e9b28b2484fbf94b543432ece638996a8b97 Mon Sep 17 00:00:00 2001 From: Chris Danis Date: Fri, 21 Aug 2026 01:13:58 -0400 Subject: [PATCH 194/211] TP: enable tensor split for LFM2/LFM2MOE (#26993) Assisted-by: deepseek-v4-flash --- src/llama-arch.cpp | 2 -- src/llama-model.cpp | 4 ++++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 408954401629..c9b504c33364 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1062,8 +1062,6 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_GRANITE_HYBRID: - case LLM_ARCH_LFM2: - case LLM_ARCH_LFM2MOE: case LLM_ARCH_MINIMAX_01: case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MINIMAX_M3: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3759c86259c7..d7874e0a9270 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -487,6 +487,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ssm_out.weight"); } if (std::regex_match(tensor_name, pattern_r_cache) || std::regex_match(tensor_name, pattern_s_cache)) { + if (ud->model->arch == LLM_ARCH_LFM2 || ud->model->arch == LLM_ARCH_LFM2MOE) { + // the LFM2 shortconv block runs fully mirrored, so its conv state must be mirrored too + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED, ""); + } return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ssm_out.weight"); } if (std::regex_match(tensor_name, pattern_ssm_conv1d)) { From 9e96cf77ffd4ebd05bad82932906ec3f59ed54ce Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Fri, 21 Aug 2026 13:14:54 +0800 Subject: [PATCH 195/211] sycl : fix load model with mlock issue (#27250) --- ggml/src/ggml-sycl/ggml-sycl.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 57aae9011d1f..83724287955e 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -1517,8 +1517,13 @@ static ggml_backend_buffer_t ggml_backend_sycl_host_buffer_type_alloc_buffer(ggm } static size_t ggml_backend_sycl_host_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - ggml_backend_sycl_device_context * dev_ctx = (ggml_backend_sycl_device_context *) buft->device->context; - return dpct::dev_mgr::instance().get_device(dev_ctx->device).get_max_mem_alloc_size(); + + if (g_ggml_sycl_enable_host_pinned_mem) { + ggml_backend_sycl_device_context * dev_ctx = (ggml_backend_sycl_device_context *) buft->device->context; + return dpct::dev_mgr::instance().get_device(dev_ctx->device).get_max_mem_alloc_size(); + } else { + return SIZE_MAX; + } } ggml_backend_buffer_type_t ggml_backend_sycl_host_buffer_type() { From 6602dd338941d1fd562e6ca3934acc5602c15a9d Mon Sep 17 00:00:00 2001 From: Ian Faust Date: Fri, 21 Aug 2026 07:15:40 +0200 Subject: [PATCH 196/211] sycl: fix multiple warnings in compiling sycl backend (#26713) * Update norm.cpp * Update helper.hpp * Update im2col.cpp * Update fattn-mkl.cpp * Update element_wise.cpp * Update fattn-mkl.cpp * Update set_rows.cpp * Update element_wise.cpp * Update ggml-sycl.cpp * Update ggml-sycl.cpp * Update ggml-sycl.cpp * Update ggml-sycl.cpp * Update ggml-sycl.cpp * Update norm.cpp * Update CMakeLists.txt * Update CMakeLists.txt * Update CMakeLists.txt * Update ggml-sycl.cpp --- ggml/src/ggml-cpu/CMakeLists.txt | 9 +++++---- ggml/src/ggml-sycl/dpct/helper.hpp | 2 +- ggml/src/ggml-sycl/element_wise.cpp | 8 ++++---- ggml/src/ggml-sycl/fattn-mkl.cpp | 9 ++++----- ggml/src/ggml-sycl/ggml-sycl.cpp | 20 +++++++++++--------- ggml/src/ggml-sycl/im2col.cpp | 4 ++-- ggml/src/ggml-sycl/norm.cpp | 8 -------- ggml/src/ggml-sycl/set_rows.cpp | 2 +- 8 files changed, 28 insertions(+), 34 deletions(-) diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt index a6cc49586bb0..32e1e7aa1932 100644 --- a/ggml/src/ggml-cpu/CMakeLists.txt +++ b/ggml/src/ggml-cpu/CMakeLists.txt @@ -737,8 +737,9 @@ function(ggml_add_cpu_backend_variant_impl tag_name) set_target_properties(${GGML_CPU_NAME} PROPERTIES COMPILE_FLAGS "-msimd128") endif() - if (CMAKE_CXX_COMPILER_ID STREQUAL "IntelLLVM") - # The compiler automatically enables "-ffast-math" which can cause NaNs in tests due to "-fassociative-math" - target_compile_options(${GGML_CPU_NAME} PRIVATE "-fno-associative-math") - endif() + if (CMAKE_C_COMPILER_ID STREQUAL "IntelLLVM" OR CMAKE_CXX_COMPILER_ID STREQUAL "IntelLLVM") + # The compiler automatically enables "-ffast-math" which can cause NaNs in tests due to "-fassociative-math" + target_compile_options(${GGML_CPU_NAME} PRIVATE "$<$,$>:$<$:/clang:>-fno-associative-math>") + endif() + endfunction() diff --git a/ggml/src/ggml-sycl/dpct/helper.hpp b/ggml/src/ggml-sycl/dpct/helper.hpp index 664b8e9697f8..85af4cab6810 100644 --- a/ggml/src/ggml-sycl/dpct/helper.hpp +++ b/ggml/src/ggml-sycl/dpct/helper.hpp @@ -62,7 +62,7 @@ #define DPCT_UNUSED(x) (void)(x) -inline void _abort(const char * str) { +[[noreturn]] inline void _abort(const char * str) { std::cerr << str << std::endl; std::abort(); } diff --git a/ggml/src/ggml-sycl/element_wise.cpp b/ggml/src/ggml-sycl/element_wise.cpp index 8619ed6f4b45..95914873e5a5 100644 --- a/ggml/src/ggml-sycl/element_wise.cpp +++ b/ggml/src/ggml-sycl/element_wise.cpp @@ -10,7 +10,7 @@ (ITEM.get_local_range(IDX) * ITEM.get_group(IDX) + ITEM.get_local_id(IDX)) static void acc_f32(const char * x, const char * y, float * dst, const int64_t ne, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, @@ -455,7 +455,7 @@ static void unary_mul_sycl(const T * x, const T * g, T * dst, const int64_t k, c namespace ggml_sycl_detail { static void acc_f32_sycl(const char *x, const char *y, float *dst, const int64_t n_elements, - const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, + const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, @@ -466,7 +466,7 @@ static void acc_f32_sycl(const char *x, const char *y, float *dst, sycl::range<3>(1, 1, SYCL_ACC_BLOCK_SIZE)), [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { acc_f32(x, y, dst, n_elements, - ne0, ne1, ne2, ne3, + ne0, ne1, ne2, nb00, nb01, nb02, nb03, ne10, ne11, ne12, ne13, nb10, nb11, nb12, nb13, @@ -970,7 +970,7 @@ static inline void ggml_sycl_op_acc(ggml_backend_sycl_context & ctx, ggml_tensor const int64_t offset = (int64_t) ((const int32_t *) dst->op_params)[3] / (int64_t) sizeof(float); ggml_sycl_detail::acc_f32_sycl(src0_d, src1_d, dst_d, ggml_nelements(dst), - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + dst->ne[0], dst->ne[1], dst->ne[2], src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3], diff --git a/ggml/src/ggml-sycl/fattn-mkl.cpp b/ggml/src/ggml-sycl/fattn-mkl.cpp index fc22b7bdb8c5..2d164a0840fd 100644 --- a/ggml/src/ggml-sycl/fattn-mkl.cpp +++ b/ggml/src/ggml-sycl/fattn-mkl.cpp @@ -43,7 +43,7 @@ static void mkl_fa_pack_q_fp16( dpct::queue_ptr stream, sycl::half * __restrict dst, const float * __restrict q_src, - int n_queries, int n_query_rows, int DKQ, + int n_queries, int DKQ, int gqa_ratio, int kvh_base_head, float q_scale, int64_t q_row_stride, int64_t q_head_stride, int64_t wg_size) { @@ -121,7 +121,7 @@ static void mkl_fa_online_softmax_chunk( float * __restrict VKQ_accum, int q0, int q_rows, int n_queries, int DV, int chunk_size, int chunk_start, - int kvh_head, int gqa_ratio, + int kvh_head, const sycl::half * mask_data, int64_t mask_head_stride, int64_t mask_row_stride, int mask_n_heads, float logit_softcap, int64_t wg_size) { @@ -473,7 +473,6 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * MKL_ACCUM(dequant_time_us, t_deq); // --- Resolve mask pointers --- - const sycl::half * mask_data = nullptr; int64_t mask_head_stride = 0; int64_t mask_row_stride = 0; int mask_n_heads = 0; @@ -547,7 +546,7 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * // 1. Pack all GQA Q heads into fp16 (full n_query_rows) mkl_fa_pack_q_fp16(stream, Q_head_f16_ptr, Q_batch, - n_queries, n_query_rows, DKQ, + n_queries, DKQ, gqa_ratio, kvh_base_head, q_scale, q_row_stride, q_head_stride, wg_size); @@ -605,7 +604,7 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr, q0, q_rows, n_queries, DV, this_chunk, chunk_start, - kvh_base_head, gqa_ratio, + kvh_base_head, mask_batch, mask_head_stride, mask_row_stride, mask_n_heads, logit_softcap, wg_size); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 83724287955e..7ebdce7fb903 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -921,16 +921,16 @@ ggml_backend_sycl_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, void * dev_ptr; if (use_usm_system) { - GGML_SYCL_DEBUG("[SYCL] allocating %lu Bytes with USM system\n", size); + GGML_SYCL_DEBUG("[SYCL] allocating %zu Bytes with USM system\n", size); dev_ptr = (void *)aligned_malloc_host(alignment, aligned_size); if (!dev_ptr) { - GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on host\n", __func__, size); + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on host\n", __func__, size); return nullptr; } } else { SYCL_CHECK(CHECK_TRY_ERROR(dev_ptr = (void *)ggml_sycl_malloc_device(size, *stream))); if (!dev_ptr) { - GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on device\n", __func__, size); + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on device\n", __func__, size); return nullptr; } } @@ -1177,7 +1177,7 @@ ggml_backend_sycl_split_buffer_init_tensor(ggml_backend_buffer_t buffer, SYCL_CHECK(CHECK_TRY_ERROR(buf = (char *)ggml_sycl_malloc_device(size, *stream))); if (!buf) { char err_buf[1024]; - snprintf(err_buf, 1023, "%s: can't allocate %lu Bytes of memory on device\n", __func__, size); + snprintf(err_buf, 1023, "%s: can't allocate %zu Bytes of memory on device\n", __func__, size); throw std::runtime_error(err_buf); } // set padding to 0 to avoid possible NaN values @@ -1651,7 +1651,7 @@ struct ggml_sycl_pool_leg : public ggml_sycl_pool { SYCL_CHECK(CHECK_TRY_ERROR(ptr = (void *)ggml_sycl_malloc_device(look_ahead_size, *qptr))); if (!ptr) { - GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on device/GPU\n", __func__, look_ahead_size); + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on device/GPU\n", __func__, look_ahead_size); return nullptr; } @@ -1663,7 +1663,7 @@ struct ggml_sycl_pool_leg : public ggml_sycl_pool { (uint32_t)(max_size/1024/1024), (uint32_t)(g_sycl_pool_size[id]/1024/1024), (uint32_t)(size/1024/1024)); #endif - // GGML_SYCL_DEBUG("ggml_sycl_pool_malloc_leg look_ahead_size=%lu, return %p\n", look_ahead_size, ptr); + // GGML_SYCL_DEBUG("ggml_sycl_pool_malloc_leg look_ahead_size=%zu, return %p\n", look_ahead_size, ptr); return ptr; } @@ -1843,7 +1843,7 @@ struct ggml_sycl_pool_host : public ggml_sycl_pool { SYCL_CHECK(CHECK_TRY_ERROR(ptr = (void *) sycl::malloc_host(size, *qptr))); if (!ptr) { - GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on host\n", __func__, size); + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on host\n", __func__, size); return nullptr; } pool_size += size; @@ -2779,9 +2779,9 @@ inline void ggml_sycl_op_mul_mat_sycl( const float * src1_ddf1_i = src1->type == GGML_TYPE_F32 ? (const float *) src1_ddf_i : src1_ddq_as_f32.get(); { +#if GGML_SYCL_DNNL const int64_t gemm_flops = (int64_t)row_diff * src1_ncols * ne10; const bool use_mkl_direct = gemm_flops < 256 * 256 * 256; -#if GGML_SYCL_DNNL if (g_ggml_sycl_enable_dnn && !use_mkl_direct) { DnnlGemmWrapper::row_gemm(ctx, row_diff, src1_ncols, ne10, src0_ddf_i, DnnlGemmWrapper::to_dt(), src1_ddf1_i, DnnlGemmWrapper::to_dt(), @@ -3518,7 +3518,9 @@ static void ggml_sycl_mul_mat_batched_sycl(ggml_backend_sycl_context & ctx, cons float * dst_ddf = static_cast(dst->data); const sycl::half * src1_f16 = static_cast(src1->data); +#if GGML_SYCL_DNNL const size_t type_size_src0 = ggml_type_size(src0->type); +#endif const size_t type_size_src1 = ggml_type_size(src1->type); bool is_src0_cont_2 = ggml_is_contiguous_2(src0); @@ -3535,6 +3537,7 @@ static void ggml_sycl_mul_mat_batched_sycl(ggml_backend_sycl_context & ctx, cons scope_op_debug_print scope_dbg_print(__func__, "/to_fp16_nc_sycl", dst, /*num_src=*/2, " : converting src1 to fp16"); +#if GGML_SYCL_DNNL // iterate tensor dims and find the slowest moving dim and stride int last_dim=0; int last_str=0; @@ -3554,7 +3557,6 @@ static void ggml_sycl_mul_mat_batched_sycl(ggml_backend_sycl_context & ctx, cons } } -#if GGML_SYCL_DNNL // oneDNN handles strided data and does not need overhead of ggml_get_to_fp16_nc_sycl const int64_t ne_src1 = src1->nb[last_str] * src1->ne[last_dim] / type_size_src1; src1_f16_alloc.alloc(ne_src1); diff --git a/ggml/src/ggml-sycl/im2col.cpp b/ggml/src/ggml-sycl/im2col.cpp index 7bf3584fb97e..e66616759465 100644 --- a/ggml/src/ggml-sycl/im2col.cpp +++ b/ggml/src/ggml-sycl/im2col.cpp @@ -85,7 +85,7 @@ static void im2col_sycl(const float * x, */ stream->parallel_for(sycl::nd_range<3>(block_nums * sycl::range<3>(1, 1, MIN(IC_KH_KW, SYCL_IM2COL_BLOCK_SIZE)), sycl::range<3>(1, 1, MIN(IC_KH_KW, SYCL_IM2COL_BLOCK_SIZE))), - [=](sycl::nd_item<3> item_ct1) { + [=](sycl::nd_item<3>) { im2col_kernel(x, dst, IC, IW, IH, OH, OW, KW, KH, IC_IH_IW, IH_IW, N_OH, KH_KW, IC_KH_KW, s0, s1, p0, p1, d0, d1); }); @@ -271,7 +271,7 @@ static void im2col_3d_sycl(const float * src, */ stream->parallel_for(sycl::nd_range<3>(block_nums * sycl::range<3>(1, 1, MIN(IC_KD_KH_KW, SYCL_IM2COL_BLOCK_SIZE)), sycl::range<3>(1, 1, MIN(IC_KD_KH_KW, SYCL_IM2COL_BLOCK_SIZE))), - [=](sycl::nd_item<3> item_ct1) { + [=](sycl::nd_item<3>) { im2col_3d_kernel(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, OH_OW, KD_KH_KW, ID_IH_IW, KH_KW, IH_IW, IC_ID_IH_IW, IC_KD_KH_KW, OW_KD_KH_KW, OD_OH_OW_IC_KD_KH_KW, OH_OW_IC_KD_KH_KW, OW_IC_KD_KH_KW, N_OD_OH, OD_OH, diff --git a/ggml/src/ggml-sycl/norm.cpp b/ggml/src/ggml-sycl/norm.cpp index 682a9f51ee74..f98a7a9542ca 100644 --- a/ggml/src/ggml-sycl/norm.cpp +++ b/ggml/src/ggml-sycl/norm.cpp @@ -7,9 +7,6 @@ static void norm_f32(const float* x, float* dst, const int ncols, const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, sycl::float2* s_sum, int block_size) { - const int nrows = item_ct1.get_group_range(2); - const int nchannels = item_ct1.get_group_range(1); - const int nthreads = item_ct1.get_local_range(2); const int sample = item_ct1.get_group(0); const int channel = item_ct1.get_group(1); @@ -155,9 +152,6 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, const float* mul = nullptr, const int64_t mul_stride_row = 0, const int64_t mul_stride_channel = 0, const int64_t mul_stride_sample = 0, const int mul_nrows = 0, const int mul_nchannels = 0, const int mul_nsamples = 0) { - const int nrows = item_ct1.get_group_range(2); - const int nchannels = item_ct1.get_group_range(1); - const int sample = item_ct1.get_group(0); const int channel = item_ct1.get_group(1); const int row = item_ct1.get_group(2); @@ -225,8 +219,6 @@ static void l2_norm_f32(const float * x, float * dst, const int ncols, const int64_t src_stride_sample, const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, const int block_size) { - const int nrows = item_ct1.get_group_range(2); - const int nchannels = item_ct1.get_group_range(1); const int row = item_ct1.get_group(2); const int channel = item_ct1.get_group(1); diff --git a/ggml/src/ggml-sycl/set_rows.cpp b/ggml/src/ggml-sycl/set_rows.cpp index 52a0bcb6ebad..5f8d881a2915 100644 --- a/ggml/src/ggml-sycl/set_rows.cpp +++ b/ggml/src/ggml-sycl/set_rows.cpp @@ -291,7 +291,7 @@ static void set_rows_sycl( stream->parallel_for( sycl::nd_range<1>(grid_size * block_size, block_size), - [=](sycl::nd_item<1> item_ct1) [[intel::reqd_sub_group_size(WARP_SIZE)]] { + [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { k_set_rows( src0_d, src1_d, dst_d, ne00, ne01, ne02, From 1cb3f5eb41d51fc98ac6b1d16ff199c427edd68d Mon Sep 17 00:00:00 2001 From: HumerousGorgon <31957201+HumerousGorgon@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:16:29 +0800 Subject: [PATCH 197/211] sycl: Update gate logic for Alchemist GPUs regarding OneDNN features. (#26635) * feat: updated gating logic of fattn-onednn.cpp * verified device types * Update ggml/src/ggml-sycl/fattn-onednn.cpp Accepted recommendations to add bmg_g31 arch. Co-authored-by: Neo Zhang * Improved SPDA gate, added documentation. * Added arch var to reworked gate, fixing build errors. * Fix trailing whitespaces. --------- Co-authored-by: Neo Zhang --- ggml/src/ggml-sycl/fattn-onednn.cpp | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index fd17a25d5edd..a501295192fb 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -21,14 +21,6 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (!g_ggml_sycl_fa_onednn) { return false; } - // Battlemage (Xe2) only, for now. On other Intel archs oneDNN's fused SDPA returns wrong results - // for some shapes (e.g. head_dim=64 on Arc / xe_hpg) -- an oneDNN bug tracked upstream at - // https://github.com/uxlfoundation/oneDNN/issues/5510. Remove this hardware limitation once that - // is fixed; until then non-BMG archs fall back to the existing FA kernel. - const gpu_arch arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch; - if (arch != gpu_arch::intel_gpu_bmg_g21 && arch != gpu_arch::intel_gpu_bmg_g31) { - return false; - } const ggml_tensor * Q = dst->src[0]; const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; @@ -60,6 +52,17 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { } } } + // This is the improved SPDA gate. Rather than gating Alchemist GPUs from all SPDA features, we instead target only the failing shapes. + // If the GPU being assessed isn't in the grouping below, it has full access to all SPDA shapes. Otherwise, if it's an Alchemist GPU, we block only the shapes with head sizes that fail. + // It is much easier to compare the device to a small list of failing cases than to define all the passing ones. + const gpu_arch arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch; + bool support_spda = !(arch == gpu_arch::intel_gpu_dg2_g10 || + arch == gpu_arch::intel_gpu_dg2_g11 || + arch == gpu_arch::intel_gpu_dg2_g12); + + if (!support_spda && K->ne[0] == 64) { + return false; + } // Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch: // very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on // some stacks; past the cap we fall back to the native FA kernel instead. From cd26896c19e6775b29a86908b5f049bbaec73305 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Thu, 20 Aug 2026 22:30:17 -0700 Subject: [PATCH 198/211] opencl: keep the vocab-scale K-quant lm_head on the CPU for Adreno A7X (compiler issue workaround) (#26440) * opencl: keep the vocab-scale K-quant lm_head on the CPU on the Adreno A7X * opencl: revise comments --------- Co-authored-by: Li He --- ggml/src/ggml-opencl/ggml-opencl.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 26f952a173ad..84f854cc2634 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7393,6 +7393,19 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K || op->src[0]->type == GGML_TYPE_Q6_K) { + // The E031.41 compiler (usually with A7x) miscompiles the flat K-quant + // GEMV kernels (kernel_mul_mv_q*_K_f32_flat) and makes lm_head run much + // slower than it should. So, make it fallback to CPU to preserve performance + // for this compiler series. + static const char * a7x_lmhead_env = getenv("GGML_OPENCL_A7X_LMHEAD_CPU"); + static const bool a7x_lmhead_cpu = (a7x_lmhead_env == nullptr || a7x_lmhead_env[0] != '0'); + if (a7x_lmhead_cpu && + backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X && + (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K || + op->src[0]->type == GGML_TYPE_Q6_K) && + op->src[0]->ne[1] >= 32768) { // vocab-scale weight; no FFN/attn weight is this tall + return false; + } return op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); } else if (op->src[0]->type == GGML_TYPE_Q8_0) { return op->src[1]->type == GGML_TYPE_F32; From 9e89a196b8141c8bec0ad4c0bfd0cfc73bb8cdea Mon Sep 17 00:00:00 2001 From: Todd Malsbary Date: Fri, 21 Aug 2026 00:23:02 -0700 Subject: [PATCH 199/211] sycl : Add Q5_K ESIMD kernel (#26376) * Add DMMV Q4_K and Q6_K ESIMD kernels Configure cmake build with -DGGML_SYCL_ESIMD=ON to enable. Signed-off-by: Todd Malsbary * Refactor ESIMD kernels to share common code Signed-off-by: Todd Malsbary * Move control of ESIMD from compile to runtime Signed-off-by: Todd Malsbary * Use ESIMD by default when available Signed-off-by: Todd Malsbary * Fix possible error when using ESIMD by default While not an issue in the current version, this will become an issue when additional QK ESIMD kernels are added (such as Q2_K). Signed-off-by: Todd Malsbary * Add explicit unroll to ESIMD kernels Signed-off-by: Todd Malsbary * Tidy up ESIMD kernels a bit Signed-off-by: Todd Malsbary * Add DMMV Q5_K ESIMD kernel Signed-off-by: Todd Malsbary * Remove redundant copyright notice Signed-off-by: Todd Malsbary --------- Signed-off-by: Todd Malsbary --- ggml/src/ggml-sycl/dmmv.cpp | 27 ++++++- ggml/src/ggml-sycl/esimd.hpp | 134 ++++++++++++++++++++++++++++--- ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + 3 files changed, 149 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-sycl/dmmv.cpp b/ggml/src/ggml-sycl/dmmv.cpp index d8da0a16ba92..fdcadbf91f9d 100644 --- a/ggml/src/ggml-sycl/dmmv.cpp +++ b/ggml/src/ggml-sycl/dmmv.cpp @@ -1955,6 +1955,23 @@ static void dequantize_mul_mat_vec_q4_K_sycl_reorder_esimd(const void *vx, const }); } +static void dequantize_mul_mat_vec_q5_K_sycl_reorder_esimd(const void *vx, const float *y, + float *dst, const int ncols, + const int nrows, + dpct::queue_ptr stream) { + GGML_ASSERT(ncols % QK_K == 0); + const int workgroups = (nrows + 1) / 2; + stream->submit([&](sycl::handler &h) { + sycl::local_accessor lmem(sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE * 2), h); + h.parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t)workgroups * GGML_SYCL_DMMV_ESIMD_WG_SIZE), sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE)), + [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { + dequantize_mul_mat_vec_reorder_esimd( + vx, y, dst, ncols, nrows, lmem, it); + }); + }); +} + static void dequantize_mul_mat_vec_q6_K_sycl_reorder_esimd(const void *vx, const float *y, float *dst, const int ncols, const int nrows, @@ -2134,7 +2151,15 @@ void ggml_sycl_op_dequantize_mul_mat_vec( case GGML_TYPE_Q5_K: if ((ggml_tensor_extra_gpu *) dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { - dequantize_mul_mat_vec_q5_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); +#ifdef GGML_SYCL_DMMV_HAS_ESIMD + if (g_ggml_sycl_enable_esimd) { + dequantize_mul_mat_vec_q5_K_sycl_reorder_esimd(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); + } + else +#endif + { + dequantize_mul_mat_vec_q5_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); + } } else { dequantize_mul_mat_vec_q5_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); } diff --git a/ggml/src/ggml-sycl/esimd.hpp b/ggml/src/ggml-sycl/esimd.hpp index d7609b11fec6..04e596ed3d3b 100644 --- a/ggml/src/ggml-sycl/esimd.hpp +++ b/ggml/src/ggml-sycl/esimd.hpp @@ -1,15 +1,3 @@ -// -// MIT license -// Copyright (C) 2026 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - #ifndef GGML_SYCL_ESIMD_HPP #define GGML_SYCL_ESIMD_HPP @@ -287,6 +275,128 @@ template <> struct esimd_reorder_q_traits { } }; +// --------------------------------------------------------------------------- +// Q5_K, SOA reorder layout produced by reorder_qw_q5_k: +// [qs: nb*(QK_K/2)] [qh: nb*(QK_K/8)] [scales: nb*K_SCALE_SIZE] [dm: nb*sizeof(half2)] +// with nb = nrows*num_blocks_per_row. +// +// Identical to Q4_K except each 4-bit quant gains a 5th (high) bit from qh: +// output chunk c (0..7) adds 16 when bit c of qh[l] is set, where qh[l] indexes +// the same 32 bytes for every chunk (matches dequantize_row_q5_K). +// --------------------------------------------------------------------------- +template <> struct esimd_reorder_q_traits { + struct ptrs { + const uint8_t * qs; + const uint8_t * qh; + const uint8_t * scales; + const sycl::half * dm; + }; + + static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) { + const uint8_t * qs = (const uint8_t *) vx; + const uint8_t * qh = qs + nb * (QK_K / 2); + const uint8_t * scales = qh + nb * (QK_K / 8); + const sycl::half * dm = (const sycl::half *) (scales + nb * K_SCALE_SIZE); + return { qs, qh, scales, dm }; + } + + // extract bit `bit` (0..7) of each lane and move it to bit position 4, + // e.g. for the 4-bit base quant's 5th (high) bit. `bit` is always a + // compile-time-known unrolled loop constant at call sites, so this folds + // to a single mask (bit==4), mask+left-shift (bit<4), or mask+right-shift + // (bit>4) instead of the shift+mask+shift a naive `(qh>>bit & 1) << 4` emits. + static ESIMD_INLINE sycl::ext::intel::esimd::simd extract_bit_to_pos4( + sycl::ext::intel::esimd::simd qh, int bit) { + using namespace sycl::ext::intel::esimd; + simd masked = convert(qh & simd((uint8_t) (1u << bit))); + if (bit < 4) { + return masked << simd((uint16_t) (4 - bit)); + } else if (bit > 4) { + return masked >> simd((uint16_t) (bit - 4)); + } + return masked; + } + + static ESIMD_INLINE void mac_pair( + const ptrs & pa, size_t bia, + const ptrs & pb, size_t bib, bool has_b, + sycl::ext::intel::esimd::simd & y_vec, + sycl::ext::intel::esimd::simd & acc_a, + sycl::ext::intel::esimd::simd & acc_b) { + using namespace sycl::ext::intel::esimd; + + simd qs_a = block_load(pa.qs + bia * (QK_K / 2)); + simd qs_b = 0; + simd qh_a = block_load(pa.qh + bia * (QK_K / 8)); + simd qh_b = 0; + simd scales_a = block_load(pa.scales + bia * K_SCALE_SIZE); + simd scales_b = 0; + + const float dall_a = (float) pa.dm[bia * 2 + 0]; + const float dmin_a = (float) pa.dm[bia * 2 + 1]; + float dall_b = 0.0f; + float dmin_b = 0.0f; + if (has_b) { + qs_b = block_load(pb.qs + bib * (QK_K / 2)); + qh_b = block_load(pb.qh + bib * (QK_K / 8)); + scales_b = block_load(pb.scales + bib * K_SCALE_SIZE); + dall_b = (float) pb.dm[bib * 2 + 0]; + dmin_b = (float) pb.dm[bib * 2 + 1]; + } + + simd scale_f_a, min_f_a, scale_f_b, min_f_b; + unpack_scale_min_k4(scales_a, dall_a, dmin_a, scale_f_a, min_f_a); + unpack_scale_min_k4(scales_b, dall_b, dmin_b, scale_f_b, min_f_b); + + simd qs_lo_a = qs_a & simd(0x0F); + simd qs_hi_a = qs_a >> simd(4); + simd qs_lo_b = qs_b & simd(0x0F); + simd qs_hi_b = qs_b >> simd(4); + +#pragma unroll + for (int sb = 0; sb < 8; sb += 2) { + const int q_offset = sb * 16; + simd y_lo = y_vec.select<32, 1>(sb * 32); + simd y_hi = y_vec.select<32, 1>((sb + 1) * 32); + + const float scale_a_lo = scale_f_a[sb]; + const float scale_a_hi = scale_f_a[sb + 1]; + const float min_a_lo = min_f_a[sb]; + const float min_a_hi = min_f_a[sb + 1]; + const float scale_b_lo = scale_f_b[sb]; + const float scale_b_hi = scale_f_b[sb + 1]; + const float min_b_lo = min_f_b[sb]; + const float min_b_hi = min_f_b[sb + 1]; + + simd qa_lo_u8 = qs_lo_a.select<32, 1>(q_offset); + simd qa_hi_u8 = qs_hi_a.select<32, 1>(q_offset); + simd qb_lo_u8 = qs_lo_b.select<32, 1>(q_offset); + simd qb_hi_u8 = qs_hi_b.select<32, 1>(q_offset); + simd qa_lo = convert(qa_lo_u8); + simd qa_hi = convert(qa_hi_u8); + simd qb_lo = convert(qb_lo_u8); + simd qb_hi = convert(qb_hi_u8); + + // add the 5th bit: chunk sb uses qh bit sb, chunk sb+1 uses qh bit sb+1; + // qh always indexes the same 32 bytes regardless of chunk + qa_lo += extract_bit_to_pos4(qh_a, sb); + qa_hi += extract_bit_to_pos4(qh_a, sb + 1); + qb_lo += extract_bit_to_pos4(qh_b, sb); + qb_hi += extract_bit_to_pos4(qh_b, sb + 1); + + simd deq_a_lo = convert(qa_lo) * scale_a_lo + min_a_lo; + simd deq_a_hi = convert(qa_hi) * scale_a_hi + min_a_hi; + simd deq_b_lo = convert(qb_lo) * scale_b_lo + min_b_lo; + simd deq_b_hi = convert(qb_hi) * scale_b_hi + min_b_hi; + + acc_a += y_lo * deq_a_lo; + acc_b += y_lo * deq_b_lo; + acc_a += y_hi * deq_a_hi; + acc_b += y_hi * deq_b_hi; + } + } +}; + // --------------------------------------------------------------------------- // Q6_K, SOA reorder layout: // [ql: nb*(QK_K/2)] [qh: nb*(QK_K/4)] [scales(int8): nb*(QK_K/16)] [d: nb*half] diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 7ebdce7fb903..de56ea5b9162 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -3811,6 +3811,7 @@ static bool ggml_sycl_supports_reorder_esimd(enum ggml_type type) { switch (type) { case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: return true; default: From 5fff128451d7603857597ee1fc18ac1dfb90f148 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 21 Aug 2026 10:29:17 +0300 Subject: [PATCH 200/211] test : make the FA V-is-view-of-K case a test case parameter (#27394) Resolve the TODO in test_flash_attn_ext: the branch that creates V as a sub-view of K (MLA-based models) was hardcoded for the 576/512 head shapes. Add a v_is_view_of_k test case parameter (default false) and select the sub-view branch on it; the existing 576/512 (DeepSeek MLA) cases now pass it explicitly, so the test coverage is unchanged. Also add more V-is-sub-view-of-K cases: the 320/256 (Mistral4 MLA) and 192/128 head shapes, and full views with equal head sizes (128/128 F16, 64/64 q8_0). Assisted-by: pi:llama.cpp/Qwen3.8-27B --- tests/test-backend-ops.cpp | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 89b954a7d1e0..8e3b273a1e46 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7085,9 +7085,10 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_V; std::array permute; const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) + const bool v_is_view_of_k; std::string vars() override { - return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view); + return VARS_TO_STR16(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k); } double max_nmse_err() override { @@ -7104,9 +7105,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true) + bool kv_view = true, bool v_is_view_of_k = false) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7138,14 +7139,14 @@ struct test_flash_attn_ext : public test_case { ggml_set_name(k, "k"); ggml_tensor * v = nullptr; - if (type_K == type_V && hsk_padded == 576 && hsv_padded == 512) { - // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes - - // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models + if (v_is_view_of_k) { + // the V cache is a sub-view of the K cache. this is used by some MLA-based models // for more info: // - https://github.com/ggml-org/llama.cpp/pull/13435 // - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392 // - https://github.com/ggml-org/llama.cpp/pull/18986 + GGML_ASSERT(type_K == type_V && hsv_padded <= hsk_padded); + v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0); } else { v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], kv_view); // the V tensor is usually a view of the V cache @@ -9906,12 +9907,14 @@ static std::vector> make_test_cases_eval() { if (hsk != 128 && prec == GGML_PREC_DEFAULT) continue; for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue; + // DeepSeek MLA: the V cache is a sub-view of the K cache + const bool v_is_view_of_k = hsk == 576; test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV)); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 1, 2, 3}, true, v_is_view_of_k)); // run fewer test cases permuted if (mask == true && max_bias == 0.0f && logit_softcap == 0 && kv == 512) { test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3})); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3}, true, v_is_view_of_k)); } } } @@ -9950,11 +9953,16 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - // MLA shape (V is a view of K) with quantized KV - // (the test harness builds V as a view of K for this shape; see build_graph) - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // MLA shape: the V cache is a sub-view of the K cache, with quantized KV + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes + test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 512, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). From ff14356e0caf6988f61f1f15f9dfe7d5ab398271 Mon Sep 17 00:00:00 2001 From: Todd Malsbary Date: Fri, 21 Aug 2026 01:01:40 -0700 Subject: [PATCH 201/211] sycl : add Q2_K reordered MMVQ and ESIMD kernels (#26336) * Add DMMV Q4_K and Q6_K ESIMD kernels Configure cmake build with -DGGML_SYCL_ESIMD=ON to enable. Signed-off-by: Todd Malsbary * Refactor ESIMD kernels to share common code Signed-off-by: Todd Malsbary * Move control of ESIMD from compile to runtime Signed-off-by: Todd Malsbary * Use ESIMD by default when available Signed-off-by: Todd Malsbary * Fix possible error when using ESIMD by default While not an issue in the current version, this will become an issue when additional QK ESIMD kernels are added (such as Q2_K). Signed-off-by: Todd Malsbary * Add explicit unroll to ESIMD kernels Signed-off-by: Todd Malsbary * Tidy up ESIMD kernels a bit Signed-off-by: Todd Malsbary * Add a reordered Q2_K MMVQ kernel Signed-off-by: Todd Malsbary * Add DMMV Q2_K ESIMD kernel Signed-off-by: Todd Malsbary --------- Signed-off-by: Todd Malsbary --- ggml/src/ggml-sycl/convert.cpp | 25 ++++++++- ggml/src/ggml-sycl/dequantize.hpp | 41 +++++++++++++++ ggml/src/ggml-sycl/dmmv.cpp | 27 +++++++++- ggml/src/ggml-sycl/esimd.hpp | 87 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/ggml-sycl.cpp | 2 + ggml/src/ggml-sycl/mmvq.cpp | 75 +++++++++++++++++++++++++- ggml/src/ggml-sycl/quants.hpp | 23 ++++++++ ggml/src/ggml-sycl/vecdotq.hpp | 33 ++++++++++++ 8 files changed, 309 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-sycl/convert.cpp b/ggml/src/ggml-sycl/convert.cpp index 9ec9276952dc..b660b56ab39b 100644 --- a/ggml/src/ggml-sycl/convert.cpp +++ b/ggml/src/ggml-sycl/convert.cpp @@ -76,6 +76,19 @@ static void dequantize_row_q2_K_sycl(const void *vx, dst_t *y, const int64_t k, #endif } +template +static void dequantize_row_q2_K_sycl_reorder(const void *vx, dst_t *y, const int64_t k, + dpct::queue_ptr stream) { + const int64_t nb = k / QK_K; + + dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 }); + stream->parallel_for( + sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 64), sycl::range<3>(1, 1, 64)), + [=](sycl::nd_item<3> item_ct1) { + dequantize_block_q2_K_reorder(vx, y, item_ct1, nb); + }); +} + template static void dequantize_row_q3_K_sycl(const void *vx, dst_t *y, const int64_t k, dpct::queue_ptr stream) { @@ -667,7 +680,11 @@ to_fp16_sycl_t ggml_get_to_fp16_sycl(ggml_type type, ggml_tensor * dst) { return dequantize_block_sycl; } case GGML_TYPE_Q2_K: - return dequantize_row_q2_K_sycl; + if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { + return dequantize_row_q2_K_sycl_reorder; + } else { + return dequantize_row_q2_K_sycl; + } case GGML_TYPE_Q3_K: if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { return dequantize_row_q3_K_sycl_reorder; @@ -753,7 +770,11 @@ to_fp32_sycl_t ggml_get_to_fp32_sycl(ggml_type type, ggml_tensor *dst) { return dequantize_block_sycl; } case GGML_TYPE_Q2_K: - return dequantize_row_q2_K_sycl; + if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { + return dequantize_row_q2_K_sycl_reorder; + } else { + return dequantize_row_q2_K_sycl; + } case GGML_TYPE_Q3_K: if (dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { return dequantize_row_q3_K_sycl_reorder; diff --git a/ggml/src/ggml-sycl/dequantize.hpp b/ggml/src/ggml-sycl/dequantize.hpp index 876ba1b44491..1b13e0f1a31a 100644 --- a/ggml/src/ggml-sycl/dequantize.hpp +++ b/ggml/src/ggml-sycl/dequantize.hpp @@ -943,6 +943,47 @@ static void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restri } +template +static void dequantize_block_q2_K_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, + const sycl::nd_item<3> & item_ct1, int64_t n_blocks) { +#if QK_K == 256 + const int64_t i = item_ct1.get_group(2); + if (i >= n_blocks) { + return; + } + + const uint8_t * base = static_cast(vx); + const size_t qs_offset = i * (QK_K / 4); + const size_t scales_offset = n_blocks * (QK_K / 4) + i * (QK_K / 16); + const size_t dm_offset = n_blocks * (QK_K / 4) + n_blocks * (QK_K / 16) + i * sizeof(ggml_half2); + + const uint8_t * qs = base + qs_offset; + const uint8_t * scales = base + scales_offset; + const ggml_half2 * dm = reinterpret_cast(base + dm_offset); + + const int64_t tid = item_ct1.get_local_id(2); + const int64_t n = tid / 32; + const int64_t l = tid - 32 * n; + const int64_t is = 8 * n + l / 16; + + const uint8_t q = qs[32 * n + l]; + dst_t * y = yy + i * QK_K + 128 * n; + + const float dall = (*dm)[0]; + const float dmin = (*dm)[1]; + y[l+ 0] = dall * (scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (scales[is+0] >> 4); + y[l+32] = dall * (scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (scales[is+2] >> 4); + y[l+64] = dall * (scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (scales[is+4] >> 4); + y[l+96] = dall * (scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (scales[is+6] >> 4); +#else + GGML_UNUSED(vx); + GGML_UNUSED(yy); + GGML_UNUSED(item_ct1); + GGML_UNUSED(n_blocks); + GGML_ABORT("Q2_K reorder dequantize not supported for QK_K != 256"); +#endif +} + template static void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy, const sycl::nd_item<3> &item_ct1) { diff --git a/ggml/src/ggml-sycl/dmmv.cpp b/ggml/src/ggml-sycl/dmmv.cpp index fdcadbf91f9d..d47d6831a359 100644 --- a/ggml/src/ggml-sycl/dmmv.cpp +++ b/ggml/src/ggml-sycl/dmmv.cpp @@ -1921,6 +1921,23 @@ ESIMD_INLINE void dequantize_mul_mat_vec_reorder_esimd( } } +static void dequantize_mul_mat_vec_q2_K_sycl_reorder_esimd(const void *vx, const float *y, + float *dst, const int ncols, + const int nrows, + dpct::queue_ptr stream) { + GGML_ASSERT(ncols % QK_K == 0); + const int workgroups = (nrows + 1) / 2; + stream->submit([&](sycl::handler &h) { + sycl::local_accessor lmem(sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE * 2), h); + h.parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t)workgroups * GGML_SYCL_DMMV_ESIMD_WG_SIZE), sycl::range<1>(GGML_SYCL_DMMV_ESIMD_WG_SIZE)), + [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { + dequantize_mul_mat_vec_reorder_esimd( + vx, y, dst, ncols, nrows, lmem, it); + }); + }); +} + static void dequantize_mul_mat_vec_q3_K_sycl_reorder_esimd(const void *vx, const float *y, float *dst, const int ncols, const int nrows, @@ -2111,7 +2128,15 @@ void ggml_sycl_op_dequantize_mul_mat_vec( case GGML_TYPE_Q2_K: if ((ggml_tensor_extra_gpu *) dst->src[0]->extra && ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { - dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); +#ifdef GGML_SYCL_DMMV_HAS_ESIMD + if (g_ggml_sycl_enable_esimd) { + dequantize_mul_mat_vec_q2_K_sycl_reorder_esimd(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); + } + else +#endif + { + dequantize_mul_mat_vec_q2_K_sycl_reorder(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); + } } else { dequantize_mul_mat_vec_q2_K_sycl(src0_dd_i, src1_ddf_i, dst_dd_i, ne00, row_diff, stream); } diff --git a/ggml/src/ggml-sycl/esimd.hpp b/ggml/src/ggml-sycl/esimd.hpp index 04e596ed3d3b..0485ff0ceed1 100644 --- a/ggml/src/ggml-sycl/esimd.hpp +++ b/ggml/src/ggml-sycl/esimd.hpp @@ -61,6 +61,93 @@ static ESIMD_INLINE void unpack_scale_min_k4( min_f = convert(m) * (-dmin); } +// --------------------------------------------------------------------------- +// Q2_K, SOA reorder layout produced by reorder_qw_q2_k: +// [qs: nb*(QK_K/4)] [scales: nb*(QK_K/16)] [dm: nb*sizeof(half2)] +// with nb = nrows*num_blocks_per_row. +// +// 2 bits per weight. The 8 output chunks of 32 (matching dequantize_row_q2_K) +// map to super-chunk s (0..7): byte base 32*(s/4) into the 64-byte qs array, +// bit shift 2*(s%4); the low 16 lanes use scales[2s], the high 16 use +// scales[2s+1], with dl = d*(sc & 0xF), ml = dmin*(sc >> 4), deq = dl*q - ml. +// --------------------------------------------------------------------------- +template <> struct esimd_reorder_q_traits { + struct ptrs { + const uint8_t * qs; + const uint8_t * scales; + const sycl::half * dm; + }; + + static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) { + const uint8_t * qs = (const uint8_t *) vx; + const uint8_t * scales = qs + nb * (QK_K / 4); + const sycl::half * dm = (const sycl::half *) (scales + nb * (QK_K / 16)); + return { qs, scales, dm }; + } + + static ESIMD_INLINE void mac_pair( + const ptrs & pa, size_t bia, + const ptrs & pb, size_t bib, bool has_b, + sycl::ext::intel::esimd::simd & y_vec, + sycl::ext::intel::esimd::simd & acc_a, + sycl::ext::intel::esimd::simd & acc_b) { + using namespace sycl::ext::intel::esimd; + + simd qs_a = block_load(pa.qs + bia * (QK_K / 4)); + simd qs_b = 0; + simd scales_a = block_load(pa.scales + bia * (QK_K / 16)); + simd scales_b = 0; + + const float dall_a = (float) pa.dm[bia * 2 + 0]; + const float dmin_a = (float) pa.dm[bia * 2 + 1]; + float dall_b = 0.0f; + float dmin_b = 0.0f; + if (has_b) { + qs_b = block_load(pb.qs + bib * (QK_K / 4)); + scales_b = block_load(pb.scales + bib * (QK_K / 16)); + dall_b = (float) pb.dm[bib * 2 + 0]; + dmin_b = (float) pb.dm[bib * 2 + 1]; + } + + // per-chunk scale (d * (sc & 0xF)) and min (-dmin * (sc >> 4)), all 16 codes; + // min carries the negation so the dequant epilogue adds (matches Q4_K/Q5_K) + simd scale_f_a = convert(scales_a & simd(0x0F)) * dall_a; + simd min_f_a = convert(scales_a >> simd(4)) * (-dmin_a); + simd scale_f_b = convert(scales_b & simd(0x0F)) * dall_b; + simd min_f_b = convert(scales_b >> simd(4)) * (-dmin_b); + +#pragma unroll + for (int s = 0; s < 8; ++s) { + const int byte_base = 32 * (s / 4); + const uint8_t shift = (uint8_t) (2 * (s % 4)); + simd y_s = y_vec.select<32, 1>(s * 32); + + simd qa = (qs_a.select<32, 1>(byte_base) >> shift) & simd(3); + simd qb = (qs_b.select<32, 1>(byte_base) >> shift) & simd(3); + + const float scale_a_lo = scale_f_a[2 * s + 0]; + const float scale_a_hi = scale_f_a[2 * s + 1]; + const float min_a_lo = min_f_a[2 * s + 0]; + const float min_a_hi = min_f_a[2 * s + 1]; + const float scale_b_lo = scale_f_b[2 * s + 0]; + const float scale_b_hi = scale_f_b[2 * s + 1]; + const float min_b_lo = min_f_b[2 * s + 0]; + const float min_b_hi = min_f_b[2 * s + 1]; + + simd scale_vec_a = splat_lo_hi(scale_a_lo, scale_a_hi); + simd min_vec_a = splat_lo_hi(min_a_lo, min_a_hi); + simd scale_vec_b = splat_lo_hi(scale_b_lo, scale_b_hi); + simd min_vec_b = splat_lo_hi(min_b_lo, min_b_hi); + + simd deq_a = convert(qa) * scale_vec_a + min_vec_a; + simd deq_b = convert(qb) * scale_vec_b + min_vec_b; + + acc_a += y_s * deq_a; + acc_b += y_s * deq_b; + } + } +}; + // --------------------------------------------------------------------------- // Q3_K, SOA reorder layout produced by reorder_qw_q3_k: // [qs: nb*(QK_K/4)] [hmask: nb*(QK_K/8)] [scales: nb*12] [d: nb*sizeof(half)] diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index de56ea5b9162..3f82020f41f8 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -3796,6 +3796,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) { case GGML_TYPE_Q1_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: @@ -3809,6 +3810,7 @@ inline bool ggml_sycl_supports_reorder_mmvq(enum ggml_type type) { static bool ggml_sycl_supports_reorder_esimd(enum ggml_type type) { #ifdef GGML_SYCL_DMMV_HAS_ESIMD switch (type) { + case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: diff --git a/ggml/src/ggml-sycl/mmvq.cpp b/ggml/src/ggml-sycl/mmvq.cpp index 123b2a2f0305..bfccb4b08f52 100644 --- a/ggml/src/ggml-sycl/mmvq.cpp +++ b/ggml/src/ggml-sycl/mmvq.cpp @@ -1401,6 +1401,64 @@ static void mul_mat_vec_q2_K_q8_1_sycl_switch_ncols( } } +static void reorder_mul_mat_vec_q2_k_q8_1_sycl(const void * vx, const void * vy, float * dst, const int ncols, + const int nrows, dpct::queue_ptr stream) { + GGML_ASSERT(ncols % QK_K == 0); + + // Round up to a whole number of subgroup-sized workgroups; out-of-range rows are skipped inside the kernel. + constexpr size_t num_subgroups = WARP_SIZE; + const int block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups); + const sycl::range<3> block_nums(1, 1, block_num_y); + const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE); + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims), + [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + mul_mat_vec_q_reorder>(vx, vy, dst, ncols, nrows, + nd_item); + }); + }); +} + +template +static void reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols( + const void * vx, const void * vy, float * dst, + const int ncols, const int nrows, + const int stride_col_y_bytes, const int stride_col_dst, + dpct::queue_ptr stream) { + GGML_ASSERT(ncols % QK_K == 0); + constexpr size_t num_subgroups = WARP_SIZE; + const int block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups); + const sycl::range<3> block_nums(1, 1, block_num_y); + const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE); + + stream->submit([&](sycl::handler & cgh) { + cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims), + [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + mul_mat_vec_q_reorder_ncols, ncols_dst>( + vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item); + }); + }); +} + +static void reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols( + const void * vx, const void * vy, float * dst, + const int ncols, const int nrows, const int ncols_dst, + const int stride_col_y_bytes, const int stride_col_dst, + dpct::queue_ptr stream) { + switch (ncols_dst) { + case 1: reorder_mul_mat_vec_q2_k_q8_1_sycl(vx, vy, dst, ncols, nrows, stream); break; + case 2: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<2>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 3: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<3>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 4: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<4>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 5: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<5>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 6: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<6>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 7: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<7>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + case 8: reorder_mul_mat_vec_q2_k_q8_1_sycl_ncols<8>(vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, stream); break; + default: GGML_ABORT("unsupported ncols_dst=%d for Q2_K reorder multi-col MMVQ", ncols_dst); + } +} + static void mul_mat_vec_q3_K_q8_1_sycl(const void *vx, const void *vy, float *dst, const int ncols, const int nrows, @@ -2297,7 +2355,21 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens } break; case GGML_TYPE_Q2_K: - if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) { + if ((ggml_tensor_extra_gpu *) dst->src[0]->extra && + ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder) { + if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) { + const int stride_col_y_bytes = src1_padded_col_size * q8_1_ts / q8_1_bs; + const int stride_col_dst = dst->ne[0]; + GGML_SYCL_DEBUG("Calling reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols ncols=%d\n", (int)src1_ncols); + reorder_mul_mat_vec_q2_k_q8_1_sycl_switch_ncols( + src0_dd_i, src1_ddq_i, dst_dd_i, ne00, row_diff, + src1_ncols, stride_col_y_bytes, stride_col_dst, stream); + return; + } else { + GGML_SYCL_DEBUG("Calling reorder_mul_mat_vec_q2_k_q8_1_sycl\n"); + reorder_mul_mat_vec_q2_k_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream); + } + } else if (i == 0 && src1_ncols > 1 && src1_ncols <= 8) { const int stride_col_y = src1_padded_col_size / QK8_1; const int stride_col_dst = dst->ne[0]; GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_K_q8_1_sycl_switch_ncols ncols=%d\n", (int)src1_ncols); @@ -2306,6 +2378,7 @@ void ggml_sycl_op_mul_mat_vec_q(ggml_backend_sycl_context & ctx, const ggml_tens src1_ncols, stride_col_y, stride_col_dst, stream); return; } else if (i == 0 || src1_ncols == 1) { + GGML_SYCL_DEBUG("Calling mul_mat_vec_q2_K_q8_1_sycl\n"); mul_mat_vec_q2_K_q8_1_sycl(src0_dd_i, src1_ddq_i_bs, dst_dd_i_bs, ne00, row_diff, stream); } break; diff --git a/ggml/src/ggml-sycl/quants.hpp b/ggml/src/ggml-sycl/quants.hpp index 95287f17510a..a26a6ce6e6d7 100644 --- a/ggml/src/ggml-sycl/quants.hpp +++ b/ggml/src/ggml-sycl/quants.hpp @@ -58,6 +58,29 @@ template <> struct block_q_t { static constexpr int block_to_q8_1_ratio() { return traits::qk / QK8_1; } }; +template <> struct block_q_t { + struct traits { + static constexpr uint32_t qk = QK_K; + static constexpr uint32_t qi = QI2_K; + static constexpr uint32_t qr = QR2_K; + static constexpr uint32_t vdr_mmvq = 1; + }; + + // Reordered layout: [qs (QK_K/4 per block)] [scales (QK_K/16 per block)] [dm] + static constexpr std::pair get_block_offset(const int block_index, const int /* n_blocks */) { + return { block_index * (QK_K / 4), 0 }; + } + + static constexpr std::pair get_d_offset(int nrows, int ncols, const int block_index) { + auto nblocks = (nrows * (ncols / QK_K)); + auto total_qs_bytes = nblocks * (QK_K / 4); + return { total_qs_bytes + block_index * (QK_K / 16), + total_qs_bytes + nblocks * (QK_K / 16) + block_index * sizeof(ggml_half2) }; + } + + static constexpr int block_to_q8_1_ratio() { return traits::qk / QK8_1; } +}; + template <> struct block_q_t { struct traits { static constexpr uint32_t qk = QK_K; diff --git a/ggml/src/ggml-sycl/vecdotq.hpp b/ggml/src/ggml-sycl/vecdotq.hpp index c11a6e8f9cbd..3ad4cee93a14 100644 --- a/ggml/src/ggml-sycl/vecdotq.hpp +++ b/ggml/src/ggml-sycl/vecdotq.hpp @@ -429,6 +429,39 @@ template <> struct reorder_vec_dot_q_sycl { } }; +template <> struct reorder_vec_dot_q_sycl { + static constexpr ggml_type gtype = GGML_TYPE_Q2_K; + + using q2_k_block = ggml_sycl_reordered::block_q_t; + using q2_k_traits = typename q2_k_block::traits; + + __dpct_inline__ float operator()(const void * __restrict__ vbq, const std::pair ibx_offset, + const std::pair d_offset, const int8_t * q8_1_quant_ptr, + const sycl::half2 * q8_1_ds, const int & iqs) { + const uint8_t * base = static_cast(vbq); + const uint8_t * qs = base + ibx_offset.first; + const uint8_t * scales = base + d_offset.first; + const ggml_half2 * dm = reinterpret_cast(base + d_offset.second); + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + + const int v = get_int_from_uint8_aligned(qs, iqs); + + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int8_t * quant_base_ptr = q8_1_quant_ptr + (bq8_offset + i) * QK8_1; + u[i] = get_int_from_int8_aligned(quant_base_ptr, iqs % QI8_1); + d8[i] = (*(q8_1_ds + bq8_offset + i))[0]; + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales + scale_offset, *dm, d8); + } +}; + template <> struct reorder_vec_dot_q_sycl { static constexpr ggml_type gtype = GGML_TYPE_Q3_K; From 62b22690602665a0f34ce5f722915839f4e4913d Mon Sep 17 00:00:00 2001 From: Charles Xu Date: Fri, 21 Aug 2026 10:33:30 +0200 Subject: [PATCH 202/211] kleidiai : add SME2 F32 GEMV kernel support (#26891) --- ggml/src/ggml-cpu/CMakeLists.txt | 3 ++ ggml/src/ggml-cpu/kleidiai/kernels.cpp | 48 ++++++++++++++++--------- ggml/src/ggml-cpu/kleidiai/kleidiai.cpp | 48 ++++++++++++++++++------- 3 files changed, 70 insertions(+), 29 deletions(-) diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt index 32e1e7aa1932..e16ac996a4a9 100644 --- a/ggml/src/ggml-cpu/CMakeLists.txt +++ b/ggml/src/ggml-cpu/CMakeLists.txt @@ -639,6 +639,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/ ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/ + ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/ ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/) set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}") @@ -701,6 +702,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S + ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.c + ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla_asm.S ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c diff --git a/ggml/src/ggml-cpu/kleidiai/kernels.cpp b/ggml/src/ggml-cpu/kleidiai/kernels.cpp index 3c31ab9d35f0..70b519f29ced 100644 --- a/ggml/src/ggml-cpu/kleidiai/kernels.cpp +++ b/ggml/src/ggml-cpu/kleidiai/kernels.cpp @@ -23,6 +23,7 @@ #include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h" #include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h" #include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h" +#include "kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h" #include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h" #include "kai_lhs_pack_bf16p2vlx2_f32_sme.h" @@ -76,6 +77,21 @@ static inline void kernel_run_fn10(size_t m, size_t n, size_t k, size_t /*bl*/, Fn(m, n, k, lhs, rhs, dst, dst_stride_row, dst_stride_col, clamp_min, clamp_max); } +template +static inline void kernel_run_lhs_stride_fn10(size_t m, + size_t n, + size_t k, + size_t lhs_stride, + const void * lhs, + const void * rhs, + void * dst, + size_t dst_stride_row, + size_t dst_stride_col, + float clamp_min, + float clamp_max) { + Fn(m, n, k, lhs, lhs_stride, rhs, dst, dst_stride_row, dst_stride_col, clamp_min, clamp_max); +} + template static inline void kernel_run_float_fn10(size_t m, size_t n, size_t k, size_t /*bl*/, const void* lhs, const void* rhs, void* dst, @@ -947,25 +963,25 @@ static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = { /* .packed_size_ex = */ &lhs_ps_fn5, /* .pack_func_ex = */ &lhs_pack_void_fn9, }, - /* SME GEMV */ + /* SME2 GEMV */ { - /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa, - /* .get_lhs_offset_ex = */ nullptr, - /* .get_rhs_packed_offset_ex = */ nullptr, - /* .run_kernel_ex = */ nullptr, + /* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_mr = */ kai_get_m_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla, + /* .get_lhs_offset_ex = */ &kernel_offs_fn2, + /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2, + /* .run_kernel_ex = */ &kernel_run_lhs_stride_fn10, }, /* .gemv_lhs_info = */ { - /* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme, - /* .get_packed_offset_ex = */ &lhs_offs_fn5, - /* .packed_size_ex = */ &lhs_ps_fn5, - /* .pack_func_ex = */ &lhs_pack_void_fn9, + /* .get_offset = */ nullptr, + /* .get_packed_offset_ex = */ nullptr, + /* .packed_size_ex = */ nullptr, + /* .pack_func_ex = */ nullptr, }, /* .rhs_info = */ { /* .packed_stride = */ nullptr, diff --git a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp index 2266c1689810..6729ae8422f2 100644 --- a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp +++ b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp @@ -696,6 +696,15 @@ class tensor_traits : public ggml::cpu::tensor_traits { } if (op->src[0]->type == GGML_TYPE_F32) { + ggml_kleidiai_kernels * primary = kernel_chain[0]; + kernel_info * gemv_kernel = primary ? &primary->gemv : nullptr; + if (is_gemv && op->src[1]->nb[0] == (int64_t) sizeof(float) && gemv_kernel && + gemv_kernel->get_lhs_offset_ex && gemv_kernel->get_rhs_packed_offset_ex && + gemv_kernel->run_kernel_ex && gemv_kernel->get_dst_offset) { + size = 0; + return true; + } + size_t cursor = 0; bool any_slot = false; @@ -811,15 +820,28 @@ class tensor_traits : public ggml::cpu::tensor_traits { return false; } - kernel_info * kernel = &kernels->gemm; + const size_t k = ne00; + const size_t m = ne11; + const size_t n = ne01; + const bool use_gemv = m == 1 && src1->nb[0] == (int64_t) sizeof(float) && + kernels->gemv.get_lhs_offset_ex && + kernels->gemv.get_rhs_packed_offset_ex && + kernels->gemv.run_kernel_ex && + kernels->gemv.get_dst_offset; + + kernel_info * kernel = use_gemv ? &kernels->gemv : &kernels->gemm; lhs_packing_info * lhs_info = &kernels->gemm_lhs_info; - if (!kernel || !lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex || - !lhs_info->packed_size_ex || !lhs_info->pack_func_ex || + if (!kernel || !kernel->get_lhs_offset_ex || !kernel->get_rhs_packed_offset_ex || !kernel->run_kernel_ex || !kernel->get_dst_offset) { return false; } + if (!use_gemv && (!lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex || + !lhs_info->packed_size_ex || !lhs_info->pack_func_ex)) { + return false; + } + const kleidiai_weight_header * header = kleidiai_weight_header_from_ptr(src0->data); const bool has_header = kleidiai_is_weight_header_valid(header); @@ -832,16 +854,14 @@ class tensor_traits : public ggml::cpu::tensor_traits { const int nth = params->nth > 0 ? params->nth : 1; const int ith = params->ith; - const size_t k = ne00; - const size_t m = ne11; - const size_t n = ne01; - const size_t mr = kernel->get_mr(); const size_t kr = kernel->get_kr(); const size_t sr = kernel->get_sr(); - const size_t lhs_packed_size = lhs_info->packed_size_ex(m, k, 0, mr, kr, sr); - GGML_ASSERT(lhs_packed_size <= params->wsize); + const size_t lhs_packed_size = use_gemv ? 0 : lhs_info->packed_size_ex(m, k, 0, mr, kr, sr); + if (!use_gemv) { + GGML_ASSERT(lhs_packed_size <= params->wsize); + } uint8_t * lhs_packed = static_cast(params->wdata); const size_t dst_stride = dst->nb[1]; @@ -853,7 +873,7 @@ class tensor_traits : public ggml::cpu::tensor_traits { const uint8_t * lhs_batch_base = static_cast(src1->data) + batch_idx * src1->nb[2]; uint8_t * dst_batch_base = static_cast(dst->data) + batch_idx * dst->nb[2]; - { + if (!use_gemv) { const int64_t m_roundup_mr = kai_roundup((int64_t)m, (int64_t)mr); int64_t max_threads = mr ? (m_roundup_mr / (int64_t)mr) : nth; max_threads = std::max(1, max_threads); @@ -903,15 +923,17 @@ class tensor_traits : public ggml::cpu::tensor_traits { const size_t n_to_process = std::min(chunk_cols, n - n_start); if (n_to_process > 0) { - const size_t lhs_packed_offset = lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr); + const size_t lhs_offset = use_gemv ? kernel->get_lhs_offset_ex(0, k, 0) + : lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr); const size_t rhs_packed_offset = kernel->get_rhs_packed_offset_ex(n_start, k, 0); const size_t dst_offset = kernel->get_dst_offset(0, n_start, dst_stride); - const void * lhs_ptr = lhs_packed + lhs_packed_offset; + const void * lhs_ptr = use_gemv ? lhs_batch_base + lhs_offset + : lhs_packed + lhs_offset; const void * rhs_ptr = rhs_base + rhs_packed_offset; float * dst_ptr = reinterpret_cast(dst_batch_base + dst_offset); - kernel->run_kernel_ex(m, n_to_process, k, 0, + kernel->run_kernel_ex(m, n_to_process, k, use_gemv ? src1->nb[1] : 0, lhs_ptr, rhs_ptr, dst_ptr, From 17197474510622a3b4ea7d0909d70b606f542b96 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 21 Aug 2026 11:33:40 +0300 Subject: [PATCH 203/211] ci : release clean-up (#27477) --- .github/workflows/make-release.yml | 12 ++++++------ .github/workflows/release.yml | 1 + scripts/make-release-desc.sh | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/.github/workflows/make-release.yml b/.github/workflows/make-release.yml index 78d7caa7aea5..95e0f34755a6 100644 --- a/.github/workflows/make-release.yml +++ b/.github/workflows/make-release.yml @@ -62,15 +62,15 @@ jobs: GITHUB_TOKEN: ${{ github.token }} with: tag_name: ${{ steps.checks.outputs.version }} - # TODO: remove the prerelease flag once the semantic versioning workflow is ready - # ref: https://github.com/ggml-org/ggml/discussions/1579 - prerelease: true + prerelease: false + # TODO: enrich the body of the release with more information body: | - > [!NOTE] - > Semantic versioning is still work in progress. - > More info can be found in https://github.com/ggml-org/ggml/discussions/1579 + ## Overview + + New version has been released. ${{ steps.desc.outputs.nightly }} + **More info:** [dist : releases and versioning of ggml-org projects](https://github.com/ggml-org/ggml/discussions/1579) ## ${{ steps.desc.outputs.changelog_title }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 61b2f5485d0d..a3c67604e89e 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1688,6 +1688,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} with: tag_name: ${{ steps.tag.outputs.name }} + prerelease: true body: |
diff --git a/scripts/make-release-desc.sh b/scripts/make-release-desc.sh index 100f855b3530..f1e4566e4a19 100755 --- a/scripts/make-release-desc.sh +++ b/scripts/make-release-desc.sh @@ -52,10 +52,10 @@ PREV="$( { git tag --list; echo "${VERSION}"; } \ if [[ -n "${PREV}" ]]; then CHANGELOG="$(git log --oneline "${PREV}..${RELEASE_COMMIT}")" - CHANGELOG_TITLE="Change log since ${PREV}" + CHANGELOG_TITLE="Changelog since ${PREV}" else CHANGELOG="(no previous release tag found)" - CHANGELOG_TITLE="Change log" + CHANGELOG_TITLE="Changelog" fi # Nightly release: the b* tag pointing at the release commit (|| true: no match is not an error) From e467c2ff6174835b4079d0acf123c5ca64615088 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 21 Aug 2026 13:20:44 +0300 Subject: [PATCH 204/211] ci : add nightly-tag.txt to make-release (#27485) As agreed in ggml discussion #1579, the official semver releases now include a nightly-tag.txt asset containing the tag of the corresponding nightly release (e.g. b10485). The Web UI assets are published to the HF bucket under the nightly tag, so this makes them discoverable for each official release. - make-release-desc.sh: expose the resolved nightly tag as a nightly_tag output - make-release.yml: create nightly-tag.txt from that tag, upload it as a release asset (skipped on dry-run), mention it in the release body and in the dry-run summary Assisted-by: pi:llama.cpp/Qwen3.8-27B --- .github/workflows/make-release.yml | 37 ++++++++++++++++++++++++++++++ scripts/make-release-desc.sh | 4 +++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/.github/workflows/make-release.yml b/.github/workflows/make-release.yml index 95e0f34755a6..451ec261f6e2 100644 --- a/.github/workflows/make-release.yml +++ b/.github/workflows/make-release.yml @@ -55,6 +55,20 @@ jobs: env: GITHUB_REPOSITORY: ${{ github.repository }} + - name: Create nightly-tag.txt + id: nightly_tag_file + run: | + NIGHTLY_TAG="${{ steps.desc.outputs.nightly_tag }}" + if [[ -z "${NIGHTLY_TAG}" ]]; then + echo "Warning: no nightly tag found for the release commit - nightly-tag.txt will not be created" + echo "create=false" >> "$GITHUB_OUTPUT" + exit 0 + fi + echo "${NIGHTLY_TAG}" > nightly-tag.txt + echo "create=true" >> "$GITHUB_OUTPUT" + echo "nightly-tag.txt:" + cat nightly-tag.txt + - name: Create release if: ${{ github.event.inputs.dry_run == 'false' }} uses: ggml-org/action-create-release@v1 @@ -70,18 +84,41 @@ jobs: New version has been released. ${{ steps.desc.outputs.nightly }} + + **Web UI:** the `nightly-tag.txt` asset contains the tag of the corresponding nightly release + **More info:** [dist : releases and versioning of ggml-org projects](https://github.com/ggml-org/ggml/discussions/1579) ## ${{ steps.desc.outputs.changelog_title }} ${{ steps.desc.outputs.changelog }} + - name: Upload nightly-tag.txt + if: ${{ github.event.inputs.dry_run == 'false' && steps.nightly_tag_file.outputs.create == 'true' }} + uses: actions/github-script@v8 + with: + github-token: ${{secrets.GITHUB_TOKEN}} + script: | + const fs = require('fs'); + const release_id = '${{ steps.create_release.outputs.id }}'; + console.log('uploadReleaseAsset', 'nightly-tag.txt'); + await github.rest.repos.uploadReleaseAsset({ + owner: context.repo.owner, + repo: context.repo.repo, + release_id: release_id, + name: 'nightly-tag.txt', + data: await fs.readFileSync('./nightly-tag.txt') + }); + - name: Dry run summary if: ${{ github.event.inputs.dry_run == 'true' }} run: | if [[ "${{ steps.checks.outputs.checks_passed }}" == "true" ]]; then echo "Dry run complete - all checks passed." echo "Would have created tag: ${{ steps.checks.outputs.version }}" + if [[ -n "${{ steps.desc.outputs.nightly_tag }}" ]]; then + echo "Would have uploaded nightly-tag.txt: ${{ steps.desc.outputs.nightly_tag }}" + fi else echo "::error::Dry run found release check failures. A release tag would not be created." exit 1 diff --git a/scripts/make-release-desc.sh b/scripts/make-release-desc.sh index f1e4566e4a19..59aa67cba7e9 100755 --- a/scripts/make-release-desc.sh +++ b/scripts/make-release-desc.sh @@ -15,7 +15,8 @@ # tag exists. # # Env (when running in GitHub Actions): -# GITHUB_OUTPUT: previous_tag, changelog_title, changelog and nightly are written here +# GITHUB_OUTPUT: previous_tag, changelog_title, changelog, nightly and nightly_tag +# are written here # GITHUB_REPOSITORY: owner/repo, used to build the nightly release URL (skipped when unset) set -euo pipefail @@ -80,6 +81,7 @@ if [[ -n "${GITHUB_OUTPUT:-}" ]]; then echo "previous_tag=${PREV}" echo "changelog_title=${CHANGELOG_TITLE}" echo "nightly=${NIGHTLY}" + echo "nightly_tag=${NIGHTLY_TAG}" echo "changelog< Date: Fri, 21 Aug 2026 12:30:03 +0200 Subject: [PATCH 205/211] ui: Settings navigation cleanup (#27241) * ui : rework the settings registry into ordered raw-data sections SETTINGS_REGISTRY becomes an ordered SettingsSectionEntry[] array; the array order is the sidebar display order. Section titles, color mode options and title radio options are declared inline in their section or entry. Entries gain showInUi; MCP servers, the system-message toggle and the title LLM flag become hidden entries of their own section. Derived values (config defaults, help info, chat sections, numeric field lists, syncable parameters) are still derived here; they move to their actual consumers in follow-up commits. * ui : extract settings localStorage persistence into SettingsService Stateless load/save of the settings config and user-override keys, plus the legacy theme key migration. Business logic (default merging, mobile sendOnEnter default, applying the migrated theme) stays in the store. * ui : move the settings exit route into ROUTES SETTINGS_FALLBACK_EXIT_ROUTE is just a route, so it lives with the other routes as ROUTES.SETTINGS_EXIT. * ui : derive the syncable parameter list in the parameter sync service The syncable parameter mapping is only consumed by the sync service, so derive it there from the registry instead of exporting it from the constants file. * ui : restore isPrivate for API key masking * ui : clean up settings registry and router fetch guard Drop the per-entry section field (duplicates the parent slug and is never read) and guard the router model fetch on fields?.length so the Tools/Import-Export pages with empty fields are excluded again. Assisted-by: pi * ui : merge sampling and penalties settings into one section Assisted-by: pi --- .../settings/SettingsChat/SettingsChat.svelte | 10 +- tools/ui/src/lib/constants/index.ts | 2 +- .../ui/src/lib/constants/routes.constants.ts | 14 +- ...try.constants.ts => settings.constants.ts} | 761 ++++++++---------- tools/ui/src/lib/services/index.ts | 10 + .../lib/services/parameter-sync.service.ts | 16 +- tools/ui/src/lib/services/settings.service.ts | 76 ++ .../src/lib/stores/settings/index.svelte.ts | 76 +- .../lib/stores/settings/referrer.svelte.ts | 4 +- tools/ui/src/lib/types/settings.d.ts | 4 +- tools/ui/src/routes/settings/+layout.svelte | 4 +- 11 files changed, 488 insertions(+), 489 deletions(-) rename tools/ui/src/lib/constants/{settings-registry.constants.ts => settings.constants.ts} (74%) create mode 100644 tools/ui/src/lib/services/settings.service.ts diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte index 4233039eff44..97ff30ba7a47 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte @@ -15,7 +15,7 @@ NUMERIC_FIELDS, POSITIVE_INTEGER_FIELDS, SETTINGS_CHAT_SECTIONS, - SETTINGS_SECTION_TITLES + SETTINGS_SECTION_SLUGS } from '$lib/constants'; import { ColorMode } from '$lib/enums/ui.enums'; import { RouterService } from '$lib/services/router.service'; @@ -46,7 +46,7 @@ let fetchInitiated = false; $effect(() => { - if (serverStore.isRouterMode && currentSection.fields && !fetchInitiated) { + if (serverStore.isRouterMode && currentSection.fields?.length && !fetchInitiated) { fetchInitiated = true; void modelsStore @@ -148,9 +148,9 @@

{currentSection.title}

- {#if currentSection.title === SETTINGS_SECTION_TITLES.TOOLS} + {#if currentSection.slug === SETTINGS_SECTION_SLUGS.TOOLS} - {:else if currentSection.title === SETTINGS_SECTION_TITLES.IMPORT_EXPORT} + {:else if currentSection.slug === SETTINGS_SECTION_SLUGS.IMPORT_EXPORT} {:else if currentSection.fields}
@@ -161,7 +161,7 @@ onThemeChange={handleThemeChange} /> - {#if currentSection.title === SETTINGS_SECTION_TITLES.GENERAL} + {#if currentSection.slug === SETTINGS_SECTION_SLUGS.GENERAL}