ui: Add extensions framework to Android Input Lifecycle plugin

Introduce the concept of extensions to allow registering custom input
lifecycle stages. This allows the plugin to dynamically fetch stages and
join them to the lifecycle from registered extensions and integrate them
chronologically into the UI.

Bug: 454761692
diff --git a/ui/src/plugins/com.android.AndroidInputLifecycle/android_input_event_source.ts b/ui/src/plugins/com.android.AndroidInputLifecycle/android_input_event_source.ts
index 7edae59..bebbb3f 100644
--- a/ui/src/plugins/com.android.AndroidInputLifecycle/android_input_event_source.ts
+++ b/ui/src/plugins/com.android.AndroidInputLifecycle/android_input_event_source.ts
@@ -17,159 +17,192 @@
   LONG_NULL,
   NUM_NULL,
   STR_NULL,
+  type Row,
+  type SqlValue,
 } from '../../trace_processor/query_result';
-import {type duration, type time, Time, Duration} from '../../base/time';
+import {type duration, Time, Duration} from '../../base/time';
 import {
   getTrackUriForTrackId,
   enrichDepths,
 } from '../../components/related_events/utils';
 import {QuerySlot, type QueryResult} from '../../base/query_slot';
-
-export interface NavTarget {
-  id: number;
-  trackUri: string;
-  ts: time;
-  dur: duration;
-  depth: number;
-}
+import type {
+  InputLifecycleExtension,
+  CellData,
+  NavTarget,
+  StageDefinition,
+} from './extensions/interface';
 
 export interface InputChainRow {
   uiRowId: string;
+  inputEventId: string | null;
   channel: string;
   totalLatency: duration | null;
-
-  durReader: duration | null;
-  durDispatch: duration | null;
-  durReceive: duration | null;
-  durConsume: duration | null;
-  durFrame: duration | null;
-
-  navReader?: NavTarget;
-  navDispatch?: NavTarget;
-  navConsume?: NavTarget;
-  navReceive?: NavTarget;
-  navFrame?: NavTarget;
-
+  stagesData: Map<string, CellData>;
   allTrackUris: string[];
 }
 
+interface InputLifecycleSpec extends Row {
+  input_id: string | null;
+  channel: string | null;
+  total_latency: bigint | null;
+  [key: string]: SqlValue;
+}
+
 export class AndroidInputEventSource {
   private readonly dataSlot = new QuerySlot<InputChainRow[]>();
 
-  constructor(private readonly trace: Trace) {}
+  constructor(
+    private readonly trace: Trace,
+    private readonly activeExtensions: ReadonlyArray<InputLifecycleExtension>,
+  ) {}
 
   use(sliceId: number): QueryResult<InputChainRow[]> {
     return this.dataSlot.use({
       key: {sliceId},
       queryFn: async () => {
-        const rows = await this.fetchRows(sliceId);
+        const resolvedSliceId = await this.resolveSliceId(sliceId);
+        const rows = await this.fetchRows(resolvedSliceId);
         await this.enrichAllDepths(rows);
         return rows;
       },
     });
   }
 
+  /**
+   * Resolves an extension-specific slice ID back to a core slice ID.
+   *
+   * If the clicked slice belongs to an extension, we resolve it to the framework
+   * input event ID, and then look up the corresponding 'InputReader' slice ID.
+   * This core slice ID is used to "root" the core lifecycle query.
+   */
+  private async resolveSliceId(sliceId: number): Promise<number> {
+    for (const ext of this.activeExtensions) {
+      if (!ext.resolveInputId) continue;
+
+      const inputId = await ext.resolveInputId(this.trace, sliceId);
+      if (!inputId) continue;
+
+      // Map the framework input ID to the core InputReader notifyMotion slice ID.
+      const coreSliceResult = await this.trace.engine.query(`
+        SELECT s.id
+        FROM android_input_events e
+        JOIN slice s ON s.ts = e.read_time AND s.track_id != 0
+        WHERE e.input_event_id = '${inputId}'
+        LIMIT 1
+      `);
+      const it = coreSliceResult.iter({id: NUM_NULL});
+      if (it.valid() && it.id !== null) {
+        return it.id;
+      }
+    }
+    return sliceId;
+  }
+
+  /**
+   * Compiles the full sequence of lifecycle stage specifications by combining
+   * the core framework stages with those injected by active extensions, and
+   * sorting them chronologically based on their sequence numbers.
+   */
+  static getStageSpecs(
+    activeExtensions: ReadonlyArray<InputLifecycleExtension>,
+  ): StageDefinition[] {
+    const specs = [...CORE_STAGES];
+    for (const ext of activeExtensions) {
+      for (const stage of ext.getStages()) {
+        specs.push({
+          ...stage,
+          key: `${ext.id}-${stage.key}`,
+        });
+      }
+    }
+    return specs.sort((a, b) => a.sequenceNumber - b.sequenceNumber);
+  }
+
+  getStageSpecs(): StageDefinition[] {
+    return AndroidInputEventSource.getStageSpecs(this.activeExtensions);
+  }
+
   private async fetchRows(sliceId: number): Promise<InputChainRow[]> {
-    const result = await this.trace.engine.query(
-      `SELECT * FROM _android_input_lifecycle_by_slice_id(${sliceId})`,
-    );
+    const baseQuery = `SELECT * FROM _android_input_lifecycle_by_slice_id(${sliceId})`;
+
+    let selectCols = 'core.*, a_evt.event_time';
+    let joinSql = '';
+
+    for (const ext of this.activeExtensions) {
+      const spec = ext.getSqlJoinSpec();
+      const alias = spec.tableAlias ?? spec.tableName;
+      selectCols += `, ${alias}.*`;
+      joinSql += ` LEFT JOIN ${spec.tableName}${spec.tableAlias ? ` AS ${spec.tableAlias}` : ''} ON ${spec.joinOn}`;
+    }
+
+    const sql = `
+      SELECT ${selectCols}
+      FROM (${baseQuery}) core
+      LEFT JOIN android_input_events a_evt ON 
+        core.input_id = a_evt.input_event_id AND
+        core.channel = a_evt.event_channel
+      ${joinSql}
+    `;
+
+    const result = await this.trace.engine.query(sql);
 
     const rows: InputChainRow[] = [];
     let index = 0;
 
-    const it = result.iter({
+    const stages = this.getStageSpecs();
+    const spec: InputLifecycleSpec = {
       input_id: STR_NULL,
       channel: STR_NULL,
       total_latency: LONG_NULL,
+    };
+    for (const stage of stages) {
+      spec[stage.idField] = NUM_NULL;
+      spec[stage.trackField] = NUM_NULL;
+      spec[stage.tsField] = LONG_NULL;
+      spec[stage.durField] = LONG_NULL;
+    }
 
-      ts_reader: LONG_NULL,
-      id_reader: NUM_NULL,
-      track_reader: NUM_NULL,
-      dur_reader: LONG_NULL,
-
-      ts_dispatch: LONG_NULL,
-      id_dispatch: NUM_NULL,
-      track_dispatch: NUM_NULL,
-      dur_dispatch: LONG_NULL,
-
-      ts_receive: LONG_NULL,
-      id_receive: NUM_NULL,
-      track_receive: NUM_NULL,
-      dur_receive: LONG_NULL,
-
-      ts_consume: LONG_NULL,
-      id_consume: NUM_NULL,
-      track_consume: NUM_NULL,
-      dur_consume: LONG_NULL,
-
-      ts_frame: LONG_NULL,
-      id_frame: NUM_NULL,
-      track_frame: NUM_NULL,
-      dur_frame: LONG_NULL,
-    });
+    const it = result.iter(spec);
 
     while (it.valid()) {
-      const navReader = this.makeNav(
-        it.id_reader,
-        it.track_reader,
-        it.ts_reader,
-        it.dur_reader,
-      );
-      const navDispatch = this.makeNav(
-        it.id_dispatch,
-        it.track_dispatch,
-        it.ts_dispatch,
-        it.dur_dispatch,
-      );
-      const navReceive = this.makeNav(
-        it.id_receive,
-        it.track_receive,
-        it.ts_receive,
-        it.dur_receive,
-      );
-      const navConsume = this.makeNav(
-        it.id_consume,
-        it.track_consume,
-        it.ts_consume,
-        it.dur_consume,
-      );
-      const navFrame = this.makeNav(
-        it.id_frame,
-        it.track_frame,
-        it.ts_frame,
-        it.dur_frame,
-      );
-
+      const stagesData = new Map<string, CellData>();
       const allTrackUris: string[] = [];
-      for (const nav of [
-        navReader,
-        navDispatch,
-        navReceive,
-        navConsume,
-        navFrame,
-      ]) {
-        if (nav) allTrackUris.push(nav.trackUri);
+
+      for (const stage of stages) {
+        const id = it.get(stage.idField) as number | null;
+        const trackId = it.get(stage.trackField) as number | null;
+        const ts = it.get(stage.tsField) as bigint | null;
+        const dur = (it.get(stage.durField) as bigint | null) ?? 0n;
+
+        const cellDur =
+          it.get(stage.durField) !== null ? Duration.fromRaw(dur) : null;
+
+        let nav: NavTarget | undefined = undefined;
+        if (id !== null && trackId !== null && ts !== null) {
+          nav = {
+            id,
+            trackUri: getTrackUriForTrackId(this.trace, trackId),
+            ts: Time.fromRaw(ts),
+            dur: Duration.fromRaw(dur),
+            depth: 0,
+          };
+          allTrackUris.push(nav.trackUri);
+        }
+
+        stagesData.set(stage.key, {dur: cellDur, nav});
       }
 
+      // TODO(ivankc) Consider how to properly handle this in the context of extensions.
+      const totalLatency =
+        it.total_latency !== null ? Duration.fromRaw(it.total_latency) : null;
+
       rows.push({
         uiRowId: `row-${index++}`,
+        inputEventId: it.input_id,
         channel: it.channel ?? '',
-        totalLatency:
-          it.total_latency !== null ? Duration.fromRaw(it.total_latency) : null,
-        durReader:
-          it.dur_reader !== null ? Duration.fromRaw(it.dur_reader) : null,
-        durDispatch:
-          it.dur_dispatch !== null ? Duration.fromRaw(it.dur_dispatch) : null,
-        durReceive:
-          it.dur_receive !== null ? Duration.fromRaw(it.dur_receive) : null,
-        durConsume:
-          it.dur_consume !== null ? Duration.fromRaw(it.dur_consume) : null,
-        durFrame: it.dur_frame !== null ? Duration.fromRaw(it.dur_frame) : null,
-        navReader,
-        navDispatch,
-        navReceive,
-        navConsume,
-        navFrame,
+        totalLatency,
+        stagesData,
         allTrackUris,
       });
 
@@ -182,33 +215,61 @@
   private async enrichAllDepths(rows: InputChainRow[]) {
     const targets: NavTarget[] = [];
     for (const row of rows) {
-      for (const nav of [
-        row.navReader,
-        row.navDispatch,
-        row.navReceive,
-        row.navConsume,
-        row.navFrame,
-      ]) {
-        if (nav) targets.push(nav);
+      for (const stageData of row.stagesData.values()) {
+        if (stageData.nav) {
+          targets.push(stageData.nav);
+        }
       }
     }
     if (targets.length === 0) return;
     await enrichDepths(this.trace, targets);
   }
-
-  private makeNav(
-    id: number | null,
-    trackId: number | null,
-    ts: bigint | null,
-    dur: bigint | null,
-  ): NavTarget | undefined {
-    if (id === null || trackId === null || ts === null) return undefined;
-    return {
-      id,
-      trackUri: getTrackUriForTrackId(this.trace, trackId),
-      ts: Time.fromRaw(ts),
-      dur: Duration.fromRaw(dur ?? 0n),
-      depth: 0,
-    };
-  }
 }
+
+const CORE_STAGES: StageDefinition[] = [
+  {
+    key: 'read',
+    headerName: 'InputReader',
+    sequenceNumber: 1000,
+    idField: 'id_reader',
+    trackField: 'track_reader',
+    tsField: 'ts_reader',
+    durField: 'dur_reader',
+  },
+  {
+    key: 'disp',
+    headerName: 'Dispatcher',
+    sequenceNumber: 2000,
+    idField: 'id_dispatch',
+    trackField: 'track_dispatch',
+    tsField: 'ts_dispatch',
+    durField: 'dur_dispatch',
+  },
+  {
+    key: 'recv',
+    headerName: 'App Receive',
+    sequenceNumber: 3000,
+    idField: 'id_receive',
+    trackField: 'track_receive',
+    tsField: 'ts_receive',
+    durField: 'dur_receive',
+  },
+  {
+    key: 'cons',
+    headerName: 'App Consume',
+    sequenceNumber: 4000,
+    idField: 'id_consume',
+    trackField: 'track_consume',
+    tsField: 'ts_consume',
+    durField: 'dur_consume',
+  },
+  {
+    key: 'frame',
+    headerName: 'App Frame',
+    sequenceNumber: 5000,
+    idField: 'id_frame',
+    trackField: 'track_frame',
+    tsField: 'ts_frame',
+    durField: 'dur_frame',
+  },
+];
diff --git a/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/interface.ts b/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/interface.ts
new file mode 100644
index 0000000..5ab29b9
--- /dev/null
+++ b/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/interface.ts
@@ -0,0 +1,75 @@
+// Copyright (C) 2026 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import type {duration, time} from '../../../base/time';
+import type {Trace} from '../../../public/trace';
+
+export interface NavTarget {
+  id: number;
+  trackUri: string;
+  ts: time;
+  dur: duration;
+  depth: number;
+}
+
+/**
+ * Defines a single processing step (stage) in the input lifecycle.
+ * Maps to a column in the plugin UI and specifies how to extract
+ * its timestamps and navigate to its timeline slice.
+ */
+export interface StageDefinition {
+  readonly key: string;
+  readonly headerName: string;
+  readonly sequenceNumber: number;
+
+  // Mapping to SQL columns in the query result
+  readonly idField: string;
+  readonly trackField: string;
+  readonly tsField: string;
+
+  // The raw duration field representing execution time:
+  readonly durField: string;
+}
+
+export interface SqlJoinSpec {
+  readonly tableName: string;
+  readonly tableAlias?: string;
+  readonly joinOn: string;
+}
+
+export interface CellData {
+  readonly dur: duration | null;
+  readonly nav?: NavTarget;
+}
+
+/**
+ * Interface for registering input lifecycle plugin extensions.
+ *
+ * Extensions allow custom or device-specific input pipelines (e.g. hardware driver
+ * touch events or vendor-specific frameworks) to integrate their lifecycle stages
+ * into the main plugin timeline and grid view.
+ */
+export interface InputLifecycleExtension {
+  readonly id: string;
+  readonly requiredModules?: string[];
+  isEligible(trace: Trace): Promise<boolean>;
+
+  // Returns the stages this extension introduces
+  getStages(): StageDefinition[];
+
+  // Returns instructions on how to join the extension table
+  getSqlJoinSpec(): SqlJoinSpec;
+
+  resolveInputId?(trace: Trace, sliceId: number): Promise<string | undefined>;
+}
diff --git a/ui/src/plugins/com.android.AndroidInputLifecycle/index.ts b/ui/src/plugins/com.android.AndroidInputLifecycle/index.ts
index 8d2d721..0d3fd15 100644
--- a/ui/src/plugins/com.android.AndroidInputLifecycle/index.ts
+++ b/ui/src/plugins/com.android.AndroidInputLifecycle/index.ts
@@ -24,11 +24,10 @@
 import {
   AndroidInputEventSource,
   type InputChainRow,
-  type NavTarget,
 } from './android_input_event_source';
 import {AndroidInputLifecycleTab} from './tab';
-import {SLICE_TRACK_KIND} from '../../public/track_kinds';
 import type {QueryResult} from '../../base/query_slot';
+import type {InputLifecycleExtension, NavTarget} from './extensions/interface';
 
 export default class AndroidInputLifecyclePlugin implements PerfettoPlugin {
   static readonly id = 'com.android.AndroidInputLifecycle';
@@ -42,7 +41,21 @@
   async onTraceLoad(trace: Trace): Promise<void> {
     await trace.engine.query('INCLUDE PERFETTO MODULE android.input;');
 
-    const source = new AndroidInputEventSource(trace);
+    const extensions: InputLifecycleExtension[] = [];
+
+    const activeExtensions: InputLifecycleExtension[] = [];
+    for (const ext of extensions) {
+      if (await ext.isEligible(trace)) {
+        if (ext.requiredModules) {
+          for (const mod of ext.requiredModules) {
+            await trace.engine.query(`INCLUDE PERFETTO MODULE ${mod};`);
+          }
+        }
+        activeExtensions.push(ext);
+      }
+    }
+
+    const source = new AndroidInputEventSource(trace, activeExtensions);
     const pinningManager = new TrackPinningManager(trace);
 
     trace.tracks.registerOverlay(
@@ -69,6 +82,7 @@
             pinningManager,
             onToggleVisibility: (rowId) => this.toggleVisibility(rowId),
             onToggleAllVisibility: () => this.toggleAllVisibility(rows ?? []),
+            activeExtensions,
           });
         },
       },
@@ -101,12 +115,6 @@
       return {data: [], isPending: false, isFresh: true};
     }
 
-    // Only handle slice tracks to avoid false positives.
-    const track = trace.tracks.getTrack(selection.trackUri);
-    if (!track?.tags?.kinds?.includes(SLICE_TRACK_KIND)) {
-      return {data: [], isPending: false, isFresh: true};
-    }
-
     return source.use(selection.eventId);
   }
 
@@ -121,13 +129,12 @@
     this.visibleRowIds.clear();
 
     for (const row of rows) {
-      const ids = [
-        row.navReader?.id,
-        row.navDispatch?.id,
-        row.navReceive?.id,
-        row.navConsume?.id,
-        row.navFrame?.id,
-      ];
+      const ids: number[] = [];
+      for (const stageData of row.stagesData.values()) {
+        if (stageData.nav) {
+          ids.push(stageData.nav.id);
+        }
+      }
       if (ids.includes(eventId)) {
         this.visibleRowIds.add(row.uiRowId);
         break;
@@ -142,17 +149,19 @@
     const {data: rows} = this.useRowState(trace, source);
     if (!rows) return [];
 
+    const specs = source.getStageSpecs();
+
     const connections: ArrowConnection[] = [];
     for (const row of rows) {
       if (!this.visibleRowIds.has(row.uiRowId)) continue;
 
-      const steps = [
-        row.navReader,
-        row.navDispatch,
-        row.navReceive,
-        row.navConsume,
-        row.navFrame,
-      ].filter((s): s is NavTarget => s !== undefined);
+      const steps: NavTarget[] = [];
+      for (const spec of specs) {
+        const stageData = row.stagesData.get(spec.key);
+        if (stageData?.nav) {
+          steps.push(stageData.nav);
+        }
+      }
 
       for (let i = 0; i < steps.length - 1; i++) {
         const start = steps[i];
diff --git a/ui/src/plugins/com.android.AndroidInputLifecycle/tab.ts b/ui/src/plugins/com.android.AndroidInputLifecycle/tab.ts
index 98d6101..4891f65 100644
--- a/ui/src/plugins/com.android.AndroidInputLifecycle/tab.ts
+++ b/ui/src/plugins/com.android.AndroidInputLifecycle/tab.ts
@@ -31,7 +31,11 @@
 import {DetailsShell} from '../../widgets/details_shell';
 import {Spinner} from '../../widgets/spinner';
 
-import type {InputChainRow, NavTarget} from './android_input_event_source';
+import {
+  AndroidInputEventSource,
+  type InputChainRow,
+} from './android_input_event_source';
+import type {InputLifecycleExtension, NavTarget} from './extensions/interface';
 
 export interface AndroidInputLifecycleTabAttrs {
   trace: Trace;
@@ -41,6 +45,7 @@
   pinningManager: TrackPinningManager;
   onToggleVisibility: (rowId: string) => void;
   onToggleAllVisibility: () => void;
+  activeExtensions: InputLifecycleExtension[];
 }
 
 export class AndroidInputLifecycleTab
@@ -77,6 +82,8 @@
     const allVisible =
       rows.length > 0 && rows.every((r) => visibleRowIds.has(r.uiRowId));
 
+    const specs = AndroidInputEventSource.getStageSpecs(attrs.activeExtensions);
+
     const columns: GridColumn[] = [
       {
         key: 'show',
@@ -95,72 +102,59 @@
         widthPx: 40,
         header: m(GridHeaderCell, {}, 'Pin'),
       },
-      {key: 'chan', header: m(GridHeaderCell, {}, 'Channel')},
+      {
+        key: 'chan',
+        header: m(GridHeaderCell, {}, 'Channel'),
+      },
       {
         key: 'total',
         minWidthPx: 100,
         header: m(GridHeaderCell, {}, 'Total Latency'),
       },
-      {
-        key: 'read',
+      ...specs.map((spec) => ({
+        key: spec.key,
         minWidthPx: 100,
-        header: m(GridHeaderCell, {}, 'InputReader'),
-      },
-      {
-        key: 'disp',
-        minWidthPx: 100,
-        header: m(GridHeaderCell, {}, 'Dispatcher'),
-      },
-      {
-        key: 'recv',
-        minWidthPx: 100,
-        header: m(GridHeaderCell, {}, 'App Receive'),
-      },
-      {
-        key: 'cons',
-        minWidthPx: 100,
-        header: m(GridHeaderCell, {}, 'App Consume'),
-      },
-      {
-        key: 'frame',
-        minWidthPx: 100,
-        header: m(GridHeaderCell, {}, 'App Frame'),
-      },
+        header: m(GridHeaderCell, {}, spec.headerName),
+      })),
     ];
 
     return m(Grid, {
       columns,
-      rowData: rows.map((row) => [
-        m(
-          GridCell,
-          {},
-          m(Checkbox, {
-            checked: visibleRowIds.has(row.uiRowId),
-            onchange: () => attrs.onToggleVisibility(row.uiRowId),
-          }),
-        ),
-        m(
-          GridCell,
-          {},
-          m(Checkbox, {
-            checked: isRowPinned(row, pinningManager),
-            onchange: () => togglePinning(row, pinningManager),
-          }),
-        ),
-        m(GridCell, {}, row.channel),
-        m(
-          GridCell,
-          {},
-          row.totalLatency !== null
-            ? m(DurationWidget, {dur: row.totalLatency, trace})
-            : '-',
-        ),
-        renderCell(row.durReader, row.navReader, trace),
-        renderCell(row.durDispatch, row.navDispatch, trace),
-        renderCell(row.durReceive, row.navReceive, trace),
-        renderCell(row.durConsume, row.navConsume, trace),
-        renderCell(row.durFrame, row.navFrame, trace),
-      ]),
+      rowData: rows.map((row) => {
+        const fixedCells = [
+          m(
+            GridCell,
+            {},
+            m(Checkbox, {
+              checked: visibleRowIds.has(row.uiRowId),
+              onchange: () => attrs.onToggleVisibility(row.uiRowId),
+            }),
+          ),
+          m(
+            GridCell,
+            {},
+            m(Checkbox, {
+              checked: isRowPinned(row, pinningManager),
+              onchange: () => togglePinning(row, pinningManager),
+            }),
+          ),
+          m(GridCell, {}, row.channel),
+          m(
+            GridCell,
+            {},
+            row.totalLatency !== null
+              ? m(DurationWidget, {dur: row.totalLatency, trace})
+              : '-',
+          ),
+        ];
+
+        const stageCells = specs.map((spec) => {
+          const cellData = row.stagesData.get(spec.key);
+          return renderCell(cellData?.dur ?? null, cellData?.nav, trace);
+        });
+
+        return [...fixedCells, ...stageCells];
+      }),
       emptyState: m(EmptyState, {
         title: 'No input event selected',
         description: 'Select an input event to see latency breakdown.',