diff --git a/packages/aiCore/src/core/providers/__tests__/HubProvider.test.ts b/packages/aiCore/src/core/providers/__tests__/HubProvider.test.ts
index 81b2050e39..ec6bf0b8cb 100644
--- a/packages/aiCore/src/core/providers/__tests__/HubProvider.test.ts
+++ b/packages/aiCore/src/core/providers/__tests__/HubProvider.test.ts
@@ -21,16 +21,11 @@ vi.mock('../core/ProviderInstanceRegistry', () => ({
   DEFAULT_SEPARATOR: '|'
 }))
 
-vi.mock('ai', async (importOriginal) => {
-  // biome-ignore lint/style/consistent-type-imports: Mock setup requires typeof import
-
-  const actual = await importOriginal<typeof AiModule>()
-  return {
-    jsonSchema: actual.jsonSchema, // Keep real jsonSchema for test utilities
-    customProvider: vi.fn((config) => config.fallbackProvider),
-    wrapProvider: vi.fn((config) => config.provider)
-  }
-})
+vi.mock('ai', () => ({
+  customProvider: vi.fn((config) => config.fallbackProvider),
+  wrapProvider: vi.fn((config) => config.provider),
+  jsonSchema: vi.fn((schema) => schema)
+}))
 
 describe('HubProvider', () => {
   let mockOpenAIProvider: ProviderV3
diff --git a/packages/aiCore/src/core/runtime/__tests__/generateImage.test.ts b/packages/aiCore/src/core/runtime/__tests__/generateImage.test.ts
index 43653f1e88..3f1b5b4231 100644
--- a/packages/aiCore/src/core/runtime/__tests__/generateImage.test.ts
+++ b/packages/aiCore/src/core/runtime/__tests__/generateImage.test.ts
@@ -10,6 +10,8 @@ import { RuntimeExecutor } from '../executor'
 // Mock dependencies
 vi.mock('ai', () => ({
   experimental_generateImage: vi.fn(),
+  generateImage: vi.fn(),
+  jsonSchema: vi.fn((schema) => schema),
   NoImageGeneratedError: class NoImageGeneratedError extends Error {
     static isInstance = vi.fn()
     constructor() {
diff --git a/packages/aiCore/src/core/types/branded.ts b/packages/aiCore/src/core/types/branded.ts
deleted file mode 100644
index 06a37ca6c8..0000000000
--- a/packages/aiCore/src/core/types/branded.ts
+++ /dev/null
@@ -1,80 +0,0 @@
-/**
- * Branded Types for type-safe IDs
- *
- * Branded types prevent accidental misuse of primitive types (like string)
- * by adding compile-time type safety without runtime overhead.
- *
- * @example
- * ```typescript
- * const modelId = ModelId('gpt-4')  // ModelId type
- * const requestId = RequestId('req-123')  // RequestId type
- *
- * function processModel(id: ModelId) { ... }
- * processModel(requestId)  // ❌ Compile error - type mismatch
- * ```
- */
-
-/**
- * Brand helper type
- */
-type Brand<K, T> = K & { readonly __brand: T }
-
-/**
- * Model ID branded type
- * Represents a unique model identifier
- */
-export type ModelId = Brand<string, 'ModelId'>
-
-/**
- * Request ID branded type
- * Represents a unique request identifier for tracing
- */
-export type RequestId = Brand<string, 'RequestId'>
-
-/**
- * Provider ID branded type
- * Represents a provider identifier (e.g., 'openai', 'anthropic')
- */
-export type ProviderId = Brand<string, 'ProviderId'>
-
-/**
- * Create a ModelId from a string
- * @param id - The model identifier string
- * @returns Branded ModelId
- */
-export const ModelId = (id: string): ModelId => id as ModelId
-
-/**
- * Create a RequestId from a string
- * @param id - The request identifier string
- * @returns Branded RequestId
- */
-export const RequestId = (id: string): RequestId => id as RequestId
-
-/**
- * Create a ProviderId from a string
- * @param id - The provider identifier string
- * @returns Branded ProviderId
- */
-export const ProviderId = (id: string): ProviderId => id as ProviderId
-
-/**
- * Type guard to check if a string is a valid ModelId
- */
-export const isModelId = (value: unknown): value is ModelId => {
-  return typeof value === 'string' && value.length > 0
-}
-
-/**
- * Type guard to check if a string is a valid RequestId
- */
-export const isRequestId = (value: unknown): value is RequestId => {
-  return typeof value === 'string' && value.length > 0
-}
-
-/**
- * Type guard to check if a string is a valid ProviderId
- */
-export const isProviderId = (value: unknown): value is ProviderId => {
-  return typeof value === 'string' && value.length > 0
-}
diff --git a/src/main/services/agents/services/claudecode/__tests__/transform.test.ts b/src/main/services/agents/services/claudecode/__tests__/transform.test.ts
index 2565f5e605..d27b426b85 100644
--- a/src/main/services/agents/services/claudecode/__tests__/transform.test.ts
+++ b/src/main/services/agents/services/claudecode/__tests__/transform.test.ts
@@ -181,7 +181,7 @@ describe('Claude → AiSDK transform', () => {
       { type: 'finish-step' }
     >
     expect(finishStep.finishReason).toBe('tool-calls')
-    expect(finishStep.usage).toEqual({ inputTokens: 1, outputTokens: 5, totalTokens: 6 })
+    expect(finishStep.usage).toMatchObject({ inputTokens: 1, outputTokens: 5, totalTokens: 6 })
 
     const toolResult = parts.find((part) => part.type === 'tool-result') as Extract<
       (typeof parts)[number],
@@ -408,7 +408,7 @@ describe('Claude → AiSDK transform', () => {
       { type: 'finish-step' }
     >
     expect(finishStep.finishReason).toBe('stop')
-    expect(finishStep.usage).toEqual({ inputTokens: 2, outputTokens: 4, totalTokens: 6 })
+    expect(finishStep.usage).toMatchObject({ inputTokens: 2, outputTokens: 4, totalTokens: 6 })
   })
 
   it('emits fallback text when Claude sends a snapshot instead of deltas', () => {
@@ -490,7 +490,7 @@ describe('Claude → AiSDK transform', () => {
       (typeof parts)[number],
       { type: 'finish-step' }
     >
-    expect(finish.usage).toEqual({ inputTokens: 3, outputTokens: 7, totalTokens: 10 })
+    expect(finish.usage).toMatchObject({ inputTokens: 3, outputTokens: 7, totalTokens: 10 })
     expect(finish.finishReason).toBe('stop')
   })
 })
diff --git a/src/main/services/agents/services/claudecode/utils.ts b/src/main/services/agents/services/claudecode/utils.ts
index eb925d7e51..14fd56ff8b 100644
--- a/src/main/services/agents/services/claudecode/utils.ts
+++ b/src/main/services/agents/services/claudecode/utils.ts
@@ -92,7 +92,7 @@ export function convertClaudeCodeUsage(usage: ClaudeCodeUsage): LanguageModelUsa
   return {
     inputTokens,
     outputTokens,
-    totalTokens: inputTokens + cacheWrite + cacheRead,
+    totalTokens: inputTokens + outputTokens,
     inputTokenDetails: {
       noCacheTokens: inputTokens,
       cacheReadTokens: cacheRead,