From c4d5f2096401d180bb20c0652f909f720c2a2f4e Mon Sep 17 00:00:00 2001
From: chrarnoldus <12196001+chrarnoldus@users.noreply.github.com>
Date: Wed, 10 Jun 2026 16:42:34 +0000
Subject: [PATCH] fix(vscode): show training badge for paid models

Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
---
 .../tests/unit/model-preview-data-line.test.ts         |  9 +++++++++
 .../webview-ui/src/components/shared/ModelPreview.tsx  |  6 ++++--
 .../webview-ui/src/components/shared/ModelSelector.tsx | 10 ++++++----
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/packages/kilo-vscode/tests/unit/model-preview-data-line.test.ts b/packages/kilo-vscode/tests/unit/model-preview-data-line.test.ts
index 63d9e27175d..788a70b61e5 100644
--- a/packages/kilo-vscode/tests/unit/model-preview-data-line.test.ts
+++ b/packages/kilo-vscode/tests/unit/model-preview-data-line.test.ts
@@ -22,6 +22,15 @@ describe("model preview data collection line", () => {
     expect(styles).toContain(".model-preview-data-line")
   })
 
+  it("renders prompt training independently from the free badge", () => {
+    expect(selector).toMatch(
+      /<Show when=\{isFree\(model\)\}>[\s\S]*?<\/Show>\s*<Show when=\{isDataCollectedModel\(model\)\}>/,
+    )
+    expect(preview).toMatch(
+      /<Show when=\{model\(\)\.isFree\}>[\s\S]*?<\/Show>\s*<Show when=\{isDataCollectedModel\(model\(\)\)\}>/,
+    )
+  })
+
   it("uses the book open check icon for all webview model data disclosures", () => {
     expect(selector).toContain('Icon name="book-open-check"')
     expect(selector).not.toContain('Icon name="warning"')
diff --git a/packages/kilo-vscode/webview-ui/src/components/shared/ModelPreview.tsx b/packages/kilo-vscode/webview-ui/src/components/shared/ModelPreview.tsx
index f4d49660bc7..ae3589f86d8 100644
--- a/packages/kilo-vscode/webview-ui/src/components/shared/ModelPreview.tsx
+++ b/packages/kilo-vscode/webview-ui/src/components/shared/ModelPreview.tsx
@@ -91,9 +91,11 @@ export const ModelPreview: Component<Props> = (props) => {
                       })()}
                     </Show>
                   </div>
-                  <Show when={model().isFree}>
+                  <Show when={model().isFree || isDataCollectedModel(model())}>
                     <span class="model-preview-free-data">
-                      <span class="model-preview-badge model-preview-badge--free">{freeLabel()}</span>
+                      <Show when={model().isFree}>
+                        <span class="model-preview-badge model-preview-badge--free">{freeLabel()}</span>
+                      </Show>
                       <Show when={isDataCollectedModel(model())}>
                         <Tooltip value={dataLabel()} placement="top">
                           <span class="model-preview-free-data-icon" aria-label={dataLabel()}>
diff --git a/packages/kilo-vscode/webview-ui/src/components/shared/ModelSelector.tsx b/packages/kilo-vscode/webview-ui/src/components/shared/ModelSelector.tsx
index e59975e7e22..00eb157021b 100644
--- a/packages/kilo-vscode/webview-ui/src/components/shared/ModelSelector.tsx
+++ b/packages/kilo-vscode/webview-ui/src/components/shared/ModelSelector.tsx
@@ -828,11 +828,13 @@ export const ModelSelectorBase: Component<ModelSelectorBaseProps> = (props) => {
                                           )
                                         })()}
                                       </span>
-                                      <Show when={isFree(model)}>
+                                      <Show when={isFree(model) || isDataCollectedModel(model)}>
                                         <span class="model-selector-free-data">
-                                          <span class="model-selector-data-badge">
-                                            <Tag data-variant="member">{freeLabel()}</Tag>
-                                          </span>
+                                          <Show when={isFree(model)}>
+                                            <span class="model-selector-data-badge">
+                                              <Tag data-variant="member">{freeLabel()}</Tag>
+                                            </span>
+                                          </Show>
                                           <Show when={isDataCollectedModel(model)}>
                                             <Tooltip value={dataLabel()} placement="top">
                                               <span class="model-selector-free-data-icon" aria-label={dataLabel()}>