feat(litert): run Tensor-compiled LiteRT models on the Pixel TPU

__tests__/harness/nativeBoundary.ts

@@ -273,6 +273,8 @@ function makeLiteRTFake(handle: FakeEmitterHandle): LiteRTFake {
     }),
     stopGeneration: jest.fn().mockResolvedValue(undefined),
     unloadModel: jest.fn().mockResolvedValue(undefined),
+    // Not a Google Tensor phone unless a test says otherwise (e.g. a Pixel 10: supported, generation 5).
+    getTpuSupport: jest.fn().mockResolvedValue({ supported: false, tensorGeneration: null, reason: 'not_tensor' }),
     getMemoryInfo: jest.fn().mockResolvedValue({
       totalRamMb: 12000,
       usedRamMb: 4000,

__tests__/integration/models/curatedLiteRTMemoryWarning.rendered.test.tsx

@@ -117,6 +117,7 @@ jest.mock('../../../src/services/hardware', () => ({
     getTotalMemoryGB: () => mockGetTotalMemoryGB(),
     getModelRecommendation: jest.fn(() => ({ maxParameters: 14, recommendedQuantization: 'Q4_K_M', recommendedModels: [], warning: undefined })),
     getImageModelRecommendation: jest.fn(() => Promise.resolve({ recommendedBackend: 'mnn', maxModelSizeMB: 2048, canRunSD: true, canRunQNN: false })),
+    getTensorTpuGeneration: jest.fn(() => Promise.resolve(null)), // not a Pixel Tensor TPU
   },
 }));
 

__tests__/integration/models/tensorTpuLiteRTOnboarding.rendered.test.tsx

@@ -0,0 +1,43 @@
+/**
+ * Pixel 10 (Tensor G5): its TPU runs only LiteRT builds compiled for it, and its GPU has no working
+ * LiteRT path, so onboarding offers the Tensor G5 build there — and nowhere else, since no other phone
+ * can run it.
+ *
+ * Real AdvancedSetupScreen + real hardwareService + curated registry; fakes ONLY the native boundary
+ * (the LiteRT module's TPU probe answers like a Pixel 10, or like any other phone).
+ */
+import { installNativeBoundary, requireRTL, GB } from '../../harness/nativeBoundary';
+
+jest.mock('@react-navigation/native', () => ({
+  useNavigation: () => ({ navigate: () => {}, goBack: () => {}, setOptions: () => {}, addListener: () => () => {}, replace: () => {} }),
+  useRoute: () => ({ params: {} }),
+  useFocusEffect: () => {}, useIsFocused: () => true,
+}));
+
+async function renderOnboarding(tpu: { supported: boolean; tensorGeneration: number | null; reason: string | null }) {
+  const boundary = installNativeBoundary({ ram: { platform: 'android', totalBytes: 12 * GB, availBytes: 8 * GB } });
+  boundary.litert.module.getTpuSupport.mockResolvedValue(tpu);
+  const React = require('react');
+  const rtl = requireRTL();
+  const { hardwareService } = require('../../../src/services/hardware');
+  const { AdvancedSetupScreen } = require('../../../src/screens/ModelDownloadScreen');
+  await hardwareService.getDeviceInfo();
+  const nav: any = { navigate: () => {}, goBack: () => {}, setOptions: () => {}, addListener: () => () => {}, replace: () => {} };
+  const view = rtl.render(React.createElement(AdvancedSetupScreen, { navigation: nav }));
+  await rtl.waitFor(() => { expect(view.getByTestId('onboarding-litert-model-1')).toBeTruthy(); }, { timeout: 10000 });
+  return { rtl, view };
+}
+
+describe('Tensor TPU LiteRT build in onboarding', () => {
+  it('offers the Tensor G5 build on a Pixel 10 whose TPU LiteRT can reach', async () => {
+    const { rtl, view } = await renderOnboarding({ supported: true, tensorGeneration: 5, reason: null });
+    const card = await rtl.waitFor(() => view.getByTestId('onboarding-litert-model-2'), { timeout: 10000 });
+    expect(rtl.within(card).getByText('Gemma 4 E2B (Tensor TPU)')).toBeTruthy();
+  }, 60000);
+
+  it('does not offer it on a phone without a Tensor G5 TPU', async () => {
+    const { view } = await renderOnboarding({ supported: false, tensorGeneration: null, reason: 'not_tensor' });
+    expect(view.queryByText('Gemma 4 E2B (Tensor TPU)')).toBeNull();
+    expect(view.queryByTestId('onboarding-litert-model-2')).toBeNull();
+  }, 60000);
+});

__tests__/integration/settings/liteRTReleaseBackends.rendered.test.tsx

@@ -10,6 +10,113 @@ it('uses GPU at the native boundary when an older build saved NPU', async () =>
   ]);
 });
 
+const TENSOR_G5_FILE = '/models/gemma-4-E2B-it_Google_Tensor_G5.litertlm';
+const PIXEL_10_TPU = { supported: true, tensorGeneration: 5, reason: null };
+
+it('loads a Tensor G5 build on the TPU of a phone with a Tensor G5 TPU, whatever the saved backend', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu');
+  expect(boundary.litert.calls.loadModel).toEqual([[TENSOR_G5_FILE, 'npu', false, false, 4096]]);
+});
+
+it('keeps a portable LiteRT file off the NPU even on a phone with a Tensor TPU', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel('/models/gemma-4-E2B-it.litertlm', 'npu');
+  expect(boundary.litert.calls.loadModel).toEqual([['/models/gemma-4-E2B-it.litertlm', 'gpu', false, false, 4096]]);
+});
+
+it('does not send a Tensor G5 build to the NPU when the phone has no usable Tensor G5 TPU', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue({ supported: false, tensorGeneration: 5, reason: 'dispatch_lib_missing' });
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'cpu');
+  expect(boundary.litert.calls.loadModel).toEqual([[TENSOR_G5_FILE, 'cpu', false, false, 4096]]);
+});
+
+it('adopts the media the native engine came up with when a TPU load drops to text-only', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  boundary.litert.module.loadModel.mockResolvedValueOnce({ backend: 'npu', maxNumTokens: 4096, vision: false, audio: false });
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu', { supportsVision: true, supportsAudio: true });
+  expect(liteRTService.getActiveBackend()).toBe('npu');
+  expect(liteRTService.loadedTextOnly(TENSOR_G5_FILE)).toBe(true);
+  expect(liteRTService.supportsAudio()).toBe(false);
+});
+
+it('refuses an image for a vision model whose TPU engine came up text-only, before any generation path', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  boundary.litert.module.loadModel.mockResolvedValueOnce({ backend: 'npu', maxNumTokens: 4096, vision: false, audio: false });
+  const { liteRTService } = require('../../../src/services/litert');
+  const { localModelAcceptsImages, activeLocalTextCapabilities } = require('../../../src/services/engines');
+  const model = createDownloadedModel({ id: 'g5', engine: 'litert', filePath: TENSOR_G5_FILE, fileName: 'gemma-4-E2B-it_Google_Tensor_G5.litertlm', liteRTVision: true });
+  expect(localModelAcceptsImages(model)).toBe(true); // not loaded yet: the model's own flag decides
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu', { supportsVision: true });
+  expect(localModelAcceptsImages(model)).toBe(false);
+  expect(activeLocalTextCapabilities(model).vision).toBe(false);
+  await liteRTService.unloadModel();
+  expect(localModelAcceptsImages(model)).toBe(true);
+});
+
+it('refuses an image at the engine when the TPU load came up text-only, on the first turn after a lazy load', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  boundary.litert.module.loadModel.mockResolvedValueOnce({ backend: 'npu', maxNumTokens: 4096, vision: false, audio: false });
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu', { supportsVision: true });
+  // The tool loop's entry point (generateRaw) and the plain path (sendMessage) both refuse.
+  await expect(liteRTService.generateRaw('what is this?', { imageUris: ['file:///pic.png'] }))
+    .rejects.toThrow(/Images are not available for this model on this device/);
+  const onError = jest.fn();
+  await liteRTService.sendMessage('what is this?', { onToken: jest.fn(), onReasoning: jest.fn(), onComplete: jest.fn(), onError }, { imageUris: ['file:///pic.png'] });
+  expect(onError).toHaveBeenCalledWith(expect.objectContaining({ message: expect.stringMatching(/Images are not available/) }));
+  expect(boundary.litert.calls.sendMessageWithImages).toEqual([]);
+  // Audio on an engine that came up without it would reach the model as empty text: refused too.
+  await expect(liteRTService.generateRaw('', { audioUris: ['file:///note.wav'] }))
+    .rejects.toThrow(/Audio input is not available for this model on this device/);
+  expect(boundary.litert.calls.sendMessageWithMedia).toEqual([]);
+});
+
+it('reports a spent TPU context as context-full, so the chat compacts and retries instead of failing', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  const { liteRTService } = require('../../../src/services/litert');
+  const { contextCompactionService } = require('../../../src/services/contextCompaction');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu');
+  // LiteRT-LM's TPU executor refuses a prefill that would pass the context compiled into the model.
+  boundary.litert.scriptError('Prefill length (128) plus current step (3990) exceeds max sequence length (4096).');
+  const onError = jest.fn();
+  await liteRTService.sendMessage('hi', { onToken: jest.fn(), onReasoning: jest.fn(), onComplete: jest.fn(), onError });
+  await new Promise(resolve => setTimeout(resolve, 0));
+  expect(onError).toHaveBeenCalledTimes(1);
+  expect(contextCompactionService.isContextFullError(onError.mock.calls[0][0])).toBe(true);
+});
+
+it('rebuilds the native conversation from the compacted history when the chat retries a full context', async () => {
+  const boundary = installNativeBoundary();
+  boundary.litert.module.getTpuSupport.mockResolvedValue(PIXEL_10_TPU);
+  const { liteRTService } = require('../../../src/services/litert');
+  await liteRTService.loadModel(TENSOR_G5_FILE, 'gpu');
+  const history = [
+    { role: 'user', content: 'first question' }, { role: 'assistant', content: 'first answer' },
+    { role: 'user', content: 'second question' }, { role: 'assistant', content: 'second answer' },
+  ];
+  await liteRTService.prepareConversation('chat-1', 'sys', { history });
+  boundary.litert.scriptError('Status Code: 3. Message: Prefill length (128) plus current step (3990) exceeds max sequence length (4096).');
+  await liteRTService.sendMessage('third question', { onToken: jest.fn(), onReasoning: jest.fn(), onComplete: jest.fn(), onError: jest.fn() });
+  await new Promise(resolve => setTimeout(resolve, 0));
+  // The native conversation still holds the refused turn, so the retry must not reuse it.
+  const compacted = [{ role: 'user', content: 'Summary: two questions answered.' }, { role: 'assistant', content: 'Noted.' }];
+  await liteRTService.prepareConversation('chat-1', 'sys', { history: compacted });
+  const historyJson = boundary.litert.calls.resetConversation.at(-1)?.[5];
+  expect(historyJson).toBe(JSON.stringify(compacted));
+});
+
 it('offers only CPU and GPU in Chat and Model Settings even with a saved NPU preference', async () => {
   installNativeBoundary();
   const React = require('react');

__tests__/rntl/navigation/AppNavigator.test.tsx

@@ -95,6 +95,7 @@ jest.mock('../../../src/services/hardware', () => ({
     })),
     formatBytes: jest.fn((bytes: number) => `${(bytes / 1024 / 1024 / 1024).toFixed(1)} GB`),
     formatModelSize: jest.fn(() => '4.0 GB'),
+    getTensorTpuGeneration: jest.fn(() => Promise.resolve(null)), // not a Pixel Tensor TPU
   },
 }));
 

__tests__/unit/hooks/useChatModelActions.test.ts

@@ -44,7 +44,7 @@ jest.mock('../../../src/services/llm', () => ({
 // LiteRT boundary mocked as a dumb flag reader so engines.isModelReady/deriveEngineCapabilities
 // (which read this DIRECT module, not the barrel) resolve deterministically for the litert path.
 jest.mock('../../../src/services/litert', () => ({
-  liteRTService: { isModelLoaded: jest.fn(() => false) },
+  liteRTService: { isModelLoaded: jest.fn(() => false), loadedTextOnly: jest.fn(() => false) },
 }));
 
 // Get mock references after hoisting

__tests__/unit/hooks/useChatModelStateSync.test.ts

@@ -18,7 +18,7 @@ jest.mock('../../../src/services/llm', () => ({
   },
 }));
 jest.mock('../../../src/services/litert', () => ({
-  liteRTService: { isModelLoaded: jest.fn(() => false) },
+  liteRTService: { isModelLoaded: jest.fn(() => false), loadedTextOnly: jest.fn(() => false) },
 }));
 
 const { llmService } = require('../../../src/services/llm');

__tests__/unit/services/activeModelService.loaders.branches.test.ts

@@ -190,6 +190,20 @@ describe('doLoadLiteRTModel branches', () => {
     expect(ToastAndroid.showWithGravity).toHaveBeenCalled();
   });
 
+  it('caps the context at a curated build\'s own limit (the Tensor TPU build is compiled for 4096 tokens)', async () => {
+    mockedLiteRT.loadModel.mockResolvedValue(undefined);
+    mockedLiteRT.getActiveBackend.mockReturnValue('npu');
+    const fileName = 'gemma-4-E2B-it_Google_Tensor_G5.litertlm';
+    const ctx = liteCtx({
+      model: { id: `offgrid/litert-recommended/${fileName}`, engine: 'litert', fileName, filePath: `/models/${fileName}`, fileSize: 3113545589 },
+      store: { settings: { ...makeStore().settings, liteRTBackend: 'gpu', liteRTMaxTokens: 8192 } },
+    });
+
+    await doLoadTextModel(ctx);
+
+    expect(mockedLiteRT.loadModel).toHaveBeenCalledWith(`/models/${fileName}`, 'gpu', expect.objectContaining({ maxNumTokens: 4096 }));
+  });
+
   it('runs warmup when active backend is gpu/npu', async () => {
     mockedLiteRT.loadModel.mockResolvedValue(undefined);
     mockedLiteRT.getActiveBackend.mockReturnValue('npu');

__tests__/unit/services/contextCompaction.test.ts

@@ -73,6 +73,11 @@ describe('isContextFullError', () => {
     ['context window exceeded', true],
     ['context length exceeded', true],
     ['context is full', true],
+    // LiteRT-LM's own limit errors, as the native module forwards them
+    ['Max number of tokens reached, context window out of bounds', true],
+    ['Input token ids are too long. Exceeding the maximum number of tokens allowed: 4100 >= 4096', true],
+    ['Reached maximum number of tokens.', true],
+    ['Generation failed: Prefill length (128) plus current step (3990) exceeds max sequence length (4096).', true],
   ])('"%s" → %s', (msg, expected) => {
     const input = typeof msg === 'string' ? new Error(msg) : msg;
     expect(contextCompactionService.isContextFullError(input)).toBe(expected);

__tests__/unit/services/curatedLiteRTContextLimit.test.ts

@@ -6,11 +6,16 @@ it.each(CURATED_LITERT_ENTRIES)('matches the pinned artifact for $displayName wi
     engine: 'litert', id: `${LITERT_PARENT_ID}/${entry.fileName}`,
     fileName: entry.fileName, fileSize: entry.sizeBytes,
   });
-  expect(getCuratedLiteRTContextLimit(model)).toBe(32000);
+  expect(getCuratedLiteRTContextLimit(model)).toBe(entry.maxContextTokens);
   expect(getCuratedLiteRTContextLimit({ ...model, id: 'imported' })).toBeNull();
   expect(getCuratedLiteRTContextLimit({ ...model, fileSize: 1 })).toBeNull();
   const origin = { repoId: entry.hfRepoId, revision: entry.commitHash, path: entry.fileName };
-  expect(getCuratedLiteRTContextLimit({ ...model, id: 'transferred', origin })).toBe(32000);
+  expect(getCuratedLiteRTContextLimit({ ...model, id: 'transferred', origin })).toBe(entry.maxContextTokens);
   expect(getCuratedLiteRTContextLimit({ ...model, origin: { ...origin, revision: 'main' } })).toBeNull();
   expect(getCuratedLiteRTContextLimit({ ...model, engine: 'llama' })).toBeNull();
 });
+
+it('caps the Tensor TPU build at its compiled 4096-token KV cache', () => {
+  const entry = CURATED_LITERT_ENTRIES.find(e => e.fileName === 'gemma-4-E2B-it_Google_Tensor_G5.litertlm')!;
+  expect(entry.maxContextTokens).toBe(4096);
+});

__tests__/unit/services/curatedLiteRTTensorTpu.test.ts

@@ -0,0 +1,36 @@
+/**
+ * A LiteRT file compiled for a Google Tensor TPU runs only on that generation's TPU, so the
+ * curated list shows it only there (Models tab + onboarding share liteRTFileRunsOnDevice).
+ * The target generation comes from Google's file naming, read by liteRTTensorTarget.
+ */
+import { CURATED_LITERT_ENTRIES, liteRTFileRunsOnDevice } from '../../../src/services/curatedLiteRTRegistry';
+import { liteRTTensorTarget } from '../../../src/utils/modelHelpers';
+
+const tensorEntry = CURATED_LITERT_ENTRIES.find(e => liteRTTensorTarget(e.fileName) !== null);
+
+describe('Tensor TPU LiteRT builds', () => {
+  it('reads the Tensor generation from Google’s file naming', () => {
+    expect(liteRTTensorTarget('gemma-4-E2B-it_Google_Tensor_G5.litertlm')).toBe(5);
+    expect(liteRTTensorTarget('Gemma3-1B-IT_google_tensor_g4.litertlm')).toBe(4);
+    expect(liteRTTensorTarget('gemma-4-E2B-it.litertlm')).toBeNull();
+    expect(liteRTTensorTarget('gemma-4-E2B-it_Google_Tensor_G5.gguf')).toBeNull();
+  });
+
+  it('curates a Tensor G5 build (fixture ground-truth)', () => {
+    expect(tensorEntry?.fileName).toBe('gemma-4-E2B-it_Google_Tensor_G5.litertlm');
+  });
+
+  it('lists the Tensor G5 build only on a phone whose TPU is Tensor G5', () => {
+    const name = tensorEntry!.fileName;
+    expect(liteRTFileRunsOnDevice(name, 5)).toBe(true);
+    expect(liteRTFileRunsOnDevice(name, 4)).toBe(false);
+    expect(liteRTFileRunsOnDevice(name, null)).toBe(false);
+  });
+
+  it('lists portable builds on every phone', () => {
+    for (const entry of CURATED_LITERT_ENTRIES.filter(e => e !== tensorEntry)) {
+      expect(liteRTFileRunsOnDevice(entry.fileName, null)).toBe(true);
+      expect(liteRTFileRunsOnDevice(entry.fileName, 5)).toBe(true);
+    }
+  });
+});

__tests__/unit/services/generationToolLoop.branches.test.ts

@@ -46,6 +46,7 @@ jest.mock('../../../src/stores', () => ({
       addMessage: mockAddMessage,
       setStreamingMessage: mockSetStreamingMessage,
       setIsThinking: mockSetIsThinking,
+      conversations: [],
     }),
   },
   useRemoteServerStore: { getState: () => ({ activeServerId: null }) },
@@ -400,6 +401,21 @@ describe('runToolLoop — LiteRT loop branches', () => {
     expect(ctx.onFinalResponse).toHaveBeenCalledWith('recovered answer');
   });
 
+  it('leaves a full context (also Status Code 3) to compaction instead of dropping the tools', async () => {
+    mockedLiteRT.generateRaw.mockRejectedValue(new Error(
+      'Status Code: 3. Message: Prefill length (128) plus current step (3990) exceeds max sequence length (4096).',
+    ));
+
+    const ctx = createContext({
+      messages: [makeMessage({ role: 'user', content: 'hi' })],
+    });
+    await expect(runToolLoop(ctx)).rejects.toThrow(/exceeds max sequence length/);
+    const toolless = mockedLiteRT.prepareConversation.mock.calls.filter(
+      (c: any[]) => Array.isArray(c[2]?.tools) && c[2].tools.length === 0,
+    );
+    expect(toolless).toEqual([]);
+  });
+
   it('rethrows a non-parse error without retrying (line 427 negative branch)', async () => {
     mockedLiteRT.generateRaw.mockRejectedValue(new Error('out of memory'));