diff --git a/packages/kilo-vscode/tests/unit/data-provider-contract.test.ts b/packages/kilo-vscode/tests/unit/data-provider-contract.test.ts
index 4a254e5eb5f..3b3a966c509 100644
--- a/packages/kilo-vscode/tests/unit/data-provider-contract.test.ts
+++ b/packages/kilo-vscode/tests/unit/data-provider-contract.test.ts
@@ -24,6 +24,6 @@ describe("DataProvider props contract", () => {
 
   it("openFile is exposed on the context return value", () => {
     const src = fs.readFileSync(DATA_CONTEXT_FILE, "utf-8")
-    expect(src).toContain("openFile:")
+    expect(src).toMatch(/openFile:\s*props\.onOpenFile/)
   })
 })
diff --git a/packages/kilo-vscode/tests/unit/tool-info-contract.test.ts b/packages/kilo-vscode/tests/unit/tool-info-contract.test.ts
index d7c76c28adf..cceca3251f9 100644
--- a/packages/kilo-vscode/tests/unit/tool-info-contract.test.ts
+++ b/packages/kilo-vscode/tests/unit/tool-info-contract.test.ts
@@ -18,9 +18,9 @@ const MESSAGE_PART_FILE = path.join(MONOREPO_ROOT, "packages/ui/src/components/m
 describe("getToolInfo() return shape contract", () => {
   it("getToolInfo still exports expected shape fields", () => {
     const src = fs.readFileSync(MESSAGE_PART_FILE, "utf-8")
-    // ToolInfo type must still have these fields
-    expect(src).toMatch(/icon\s*:/)
-    expect(src).toMatch(/title\s*:/)
+    // ToolInfo type definition must still have these fields
+    expect(src).toMatch(/export type ToolInfo\s*=\s*\{[^}]*icon\s*:/s)
+    expect(src).toMatch(/export type ToolInfo\s*=\s*\{[^}]*title\s*:/s)
   })
 
   it("getToolInfo function is exported", () => {