From b9c87c14a23bf1599b393494c7700c95f895728a Mon Sep 17 00:00:00 2001 From: Michael Bleigh Date: Fri, 20 Mar 2026 06:40:10 -0700 Subject: [PATCH 01/32] feat(core): introduce AgentSession and rename stream events to agent events (#23159) --- packages/core/src/agent/agent-session.test.ts | 279 +++++++++++++++++ packages/core/src/agent/agent-session.ts | 212 +++++++++++++ packages/core/src/agent/mock.test.ts | 278 ++++++++--------- packages/core/src/agent/mock.ts | 290 ++++++++---------- packages/core/src/agent/types.ts | 48 +-- 5 files changed, 762 insertions(+), 345 deletions(-) create mode 100644 packages/core/src/agent/agent-session.test.ts create mode 100644 packages/core/src/agent/agent-session.ts diff --git a/packages/core/src/agent/agent-session.test.ts b/packages/core/src/agent/agent-session.test.ts new file mode 100644 index 0000000000..c390d719d4 --- /dev/null +++ b/packages/core/src/agent/agent-session.test.ts @@ -0,0 +1,279 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { AgentSession } from './agent-session.js'; +import { MockAgentProtocol } from './mock.js'; +import type { AgentEvent } from './types.js'; + +describe('AgentSession', () => { + it('should passthrough simple methods', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([{ type: 'message' }]); + await session.send({ update: { title: 't' } }); + // update, agent_start, message, agent_end = 4 events + expect(session.events).toHaveLength(4); + + let emitted = false; + session.subscribe(() => { + emitted = true; + }); + protocol.pushResponse([]); + await session.send({ update: { title: 't' } }); + expect(emitted).toBe(true); + + protocol.pushResponse([], { keepOpen: true }); + await session.send({ update: { title: 't' } }); + await session.abort(); + expect( + session.events.some( + (e) => + e.type === 'agent_end' && + (e as AgentEvent<'agent_end'>).reason === 'aborted', + ), + ).toBe(true); + }); + + it('should yield events via sendStream', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([ + { + type: 'message', + role: 'agent', + content: [{ type: 'text', text: 'hello' }], + }, + ]); + + const events: AgentEvent[] = []; + for await (const event of session.sendStream({ + message: [{ type: 'text', text: 'hi' }], + })) { + events.push(event); + } + + // agent_start, agent message, agent_end = 3 events (user message skipped) + expect(events).toHaveLength(3); + expect(events[0].type).toBe('agent_start'); + expect(events[1].type).toBe('message'); + expect((events[1] as AgentEvent<'message'>).role).toBe('agent'); + expect(events[2].type).toBe('agent_end'); + }); + + it('should filter events by streamId in sendStream', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([{ type: 'message' }]); + + const events: AgentEvent[] = []; + const stream = session.sendStream({ update: { title: 'foo' } }); + + for await (const event of stream) { + events.push(event); + } + + expect(events).toHaveLength(3); // agent_start, message, agent_end (update skipped) + const streamId = events[0].streamId; + expect(streamId).not.toBeNull(); + expect(events.every((e) => e.streamId === streamId)).toBe(true); + }); + + it('should handle events arriving before send() resolves', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([{ type: 'message' }]); + + const events: AgentEvent[] = []; + for await (const event of session.sendStream({ + update: { title: 'foo' }, + })) { + events.push(event); + } + + expect(events).toHaveLength(3); // agent_start, message, agent_end (update skipped) + expect(events[0].type).toBe('agent_start'); + expect(events[1].type).toBe('message'); + expect(events[2].type).toBe('agent_end'); + }); + + it('should return immediately from sendStream if streamId is null', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + // No response queued, so send() returns streamId: null + const events: AgentEvent[] = []; + for await (const event of session.sendStream({ + update: { title: 'foo' }, + })) { + events.push(event); + } + + expect(events).toHaveLength(0); + expect(protocol.events).toHaveLength(1); + expect(protocol.events[0].type).toBe('session_update'); + }); + + it('should skip events that occur before agent_start', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + // Custom emission to ensure events happen before agent_start + protocol.pushResponse([ + { + type: 'message', + role: 'agent', + content: [{ type: 'text', text: 'hello' }], + }, + ]); + + // We can't easily inject events before agent_start with MockAgentProtocol.pushResponse + // because it emits them all together. + // But we know session_update is emitted first. + + const events: AgentEvent[] = []; + for await (const event of session.sendStream({ + message: [{ type: 'text', text: 'hi' }], + })) { + events.push(event); + } + + // The session_update (from the 'hi' message) should be skipped. + expect(events.some((e) => e.type === 'session_update')).toBe(false); + expect(events[0].type).toBe('agent_start'); + }); + + describe('stream()', () => { + it('should replay events after eventId', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + // Create some events + protocol.pushResponse([{ type: 'message' }]); + await session.send({ update: { title: 't1' } }); + // Wait for events to be emitted + await new Promise((resolve) => setTimeout(resolve, 10)); + + const allEvents = session.events; + expect(allEvents.length).toBeGreaterThan(2); + const eventId = allEvents[1].id; + + const streamedEvents: AgentEvent[] = []; + for await (const event of session.stream({ eventId })) { + streamedEvents.push(event); + } + + expect(streamedEvents).toEqual(allEvents.slice(2)); + }); + + it('should replay events for streamId starting with agent_start', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([{ type: 'message' }]); + const { streamId } = await session.send({ update: { title: 't1' } }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + const allEvents = session.events; + const startEventIndex = allEvents.findIndex( + (e) => e.type === 'agent_start' && e.streamId === streamId, + ); + expect(startEventIndex).toBeGreaterThan(-1); + + const streamedEvents: AgentEvent[] = []; + for await (const event of session.stream({ streamId: streamId! })) { + streamedEvents.push(event); + } + + expect(streamedEvents).toEqual(allEvents.slice(startEventIndex)); + }); + + it('should continue listening for active stream after replay', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + // Start a stream but keep it open + protocol.pushResponse([{ type: 'message' }], { keepOpen: true }); + const { streamId } = await session.send({ update: { title: 't1' } }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + const streamedEvents: AgentEvent[] = []; + const streamPromise = (async () => { + for await (const event of session.stream({ streamId: streamId! })) { + streamedEvents.push(event); + } + })(); + + // Push more to the stream + await new Promise((resolve) => setTimeout(resolve, 20)); + protocol.pushToStream(streamId!, [{ type: 'message' }], { close: true }); + + await streamPromise; + + const allEvents = session.events; + const startEventIndex = allEvents.findIndex( + (e) => e.type === 'agent_start' && e.streamId === streamId, + ); + expect(streamedEvents).toEqual(allEvents.slice(startEventIndex)); + expect(streamedEvents.at(-1)?.type).toBe('agent_end'); + }); + + it('should follow an active stream if no options provided', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + protocol.pushResponse([{ type: 'message' }], { keepOpen: true }); + const { streamId } = await session.send({ update: { title: 't1' } }); + await new Promise((resolve) => setTimeout(resolve, 10)); + + const streamedEvents: AgentEvent[] = []; + const streamPromise = (async () => { + for await (const event of session.stream()) { + streamedEvents.push(event); + } + })(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + protocol.pushToStream(streamId!, [{ type: 'message' }], { close: true }); + await streamPromise; + + expect(streamedEvents.length).toBeGreaterThan(0); + expect(streamedEvents.at(-1)?.type).toBe('agent_end'); + }); + + it('should ONLY yield events for specific streamId even if newer streams exist', async () => { + const protocol = new MockAgentProtocol(); + const session = new AgentSession(protocol); + + // Stream 1 + protocol.pushResponse([{ type: 'message' }]); + const { streamId: streamId1 } = await session.send({ + update: { title: 's1' }, + }); + + // Stream 2 + protocol.pushResponse([{ type: 'message' }]); + const { streamId: streamId2 } = await session.send({ + update: { title: 's2' }, + }); + + await new Promise((resolve) => setTimeout(resolve, 20)); + + const streamedEvents: AgentEvent[] = []; + for await (const event of session.stream({ streamId: streamId1! })) { + streamedEvents.push(event); + } + + expect(streamedEvents.every((e) => e.streamId === streamId1)).toBe(true); + expect(streamedEvents.some((e) => e.type === 'agent_end')).toBe(true); + expect(streamedEvents.some((e) => e.streamId === streamId2)).toBe(false); + }); + }); +}); diff --git a/packages/core/src/agent/agent-session.ts b/packages/core/src/agent/agent-session.ts new file mode 100644 index 0000000000..0d9fc86bb0 --- /dev/null +++ b/packages/core/src/agent/agent-session.ts @@ -0,0 +1,212 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + AgentProtocol, + AgentSend, + AgentEvent, + Unsubscribe, +} from './types.js'; + +/** + * AgentSession is a wrapper around AgentProtocol that provides a more + * convenient API for consuming agent activity as an AsyncIterable. + */ +export class AgentSession implements AgentProtocol { + private _protocol: AgentProtocol; + + constructor(protocol: AgentProtocol) { + this._protocol = protocol; + } + + async send(payload: AgentSend): Promise<{ streamId: string | null }> { + return this._protocol.send(payload); + } + + subscribe(callback: (event: AgentEvent) => void): Unsubscribe { + return this._protocol.subscribe(callback); + } + + async abort(): Promise { + return this._protocol.abort(); + } + + get events(): AgentEvent[] { + return this._protocol.events; + } + + /** + * Sends a payload to the agent and returns an AsyncIterable that yields + * events for the resulting stream. + * + * @param payload The payload to send to the agent. + */ + async *sendStream(payload: AgentSend): AsyncIterable { + const result = await this._protocol.send(payload); + const streamId = result.streamId; + + if (streamId === null) { + return; + } + + yield* this.stream({ streamId }); + } + + /** + * Returns an AsyncIterable that yields events from the agent session, + * optionally replaying events from history or reattaching to an existing stream. + * + * @param options Options for replaying or reattaching to the event stream. + */ + async *stream( + options: { + eventId?: string; + streamId?: string; + } = {}, + ): AsyncIterable { + let resolve: (() => void) | undefined; + let next = new Promise((res) => { + resolve = res; + }); + + let eventQueue: AgentEvent[] = []; + const earlyEvents: AgentEvent[] = []; + let done = false; + let trackedStreamId = options.streamId; + let started = false; + + // 1. Subscribe early to avoid missing any events that occur during replay setup + const unsubscribe = this._protocol.subscribe((event) => { + if (done) return; + + if (!started) { + earlyEvents.push(event); + return; + } + + if (trackedStreamId && event.streamId !== trackedStreamId) return; + + // If we don't have a tracked stream yet, the first agent_start we see becomes it. + if (!trackedStreamId && event.type === 'agent_start') { + trackedStreamId = event.streamId ?? undefined; + } + + // If we still don't have a tracked stream and we aren't replaying everything (eventId), ignore. + if (!trackedStreamId && !options.eventId) return; + + eventQueue.push(event); + if ( + event.type === 'agent_end' && + event.streamId === (trackedStreamId ?? null) + ) { + done = true; + } + + const currentResolve = resolve; + next = new Promise((r) => { + resolve = r; + }); + currentResolve?.(); + }); + + try { + const currentEvents = this._protocol.events; + let replayStartIndex = -1; + + if (options.eventId) { + const index = currentEvents.findIndex((e) => e.id === options.eventId); + if (index !== -1) { + replayStartIndex = index + 1; + } + } else if (options.streamId) { + const index = currentEvents.findIndex( + (e) => e.type === 'agent_start' && e.streamId === options.streamId, + ); + if (index !== -1) { + replayStartIndex = index; + } + } + + if (replayStartIndex !== -1) { + for (let i = replayStartIndex; i < currentEvents.length; i++) { + const event = currentEvents[i]; + if (options.streamId && event.streamId !== options.streamId) continue; + + eventQueue.push(event); + if (event.type === 'agent_start' && !trackedStreamId) { + trackedStreamId = event.streamId ?? undefined; + } + if ( + event.type === 'agent_end' && + event.streamId === (trackedStreamId ?? null) + ) { + done = true; + break; + } + } + } + + if (!done && !trackedStreamId) { + // Find active stream in history + const activeStarts = currentEvents.filter( + (e) => e.type === 'agent_start', + ); + for (let i = activeStarts.length - 1; i >= 0; i--) { + const start = activeStarts[i]; + if ( + !currentEvents.some( + (e) => e.type === 'agent_end' && e.streamId === start.streamId, + ) + ) { + trackedStreamId = start.streamId ?? undefined; + break; + } + } + } + + // If we replayed to the end and no stream is active, and we were specifically + // replaying from an eventId (or we've already finished the stream we were looking for), we are done. + if (!done && !trackedStreamId && options.eventId) { + done = true; + } + + started = true; + + // Process events that arrived while we were replaying + for (const event of earlyEvents) { + if (done) break; + if (trackedStreamId && event.streamId !== trackedStreamId) continue; + if (!trackedStreamId && event.type === 'agent_start') { + trackedStreamId = event.streamId ?? undefined; + } + if (!trackedStreamId && !options.eventId) continue; + + eventQueue.push(event); + if ( + event.type === 'agent_end' && + event.streamId === (trackedStreamId ?? null) + ) { + done = true; + } + } + + while (true) { + if (eventQueue.length > 0) { + const eventsToYield = eventQueue; + eventQueue = []; + for (const event of eventsToYield) { + yield event; + } + } + + if (done) break; + await next; + } + } finally { + unsubscribe(); + } + } +} diff --git a/packages/core/src/agent/mock.test.ts b/packages/core/src/agent/mock.test.ts index 41672223a9..4f102d5dbd 100644 --- a/packages/core/src/agent/mock.test.ts +++ b/packages/core/src/agent/mock.test.ts @@ -5,12 +5,24 @@ */ import { describe, expect, it } from 'vitest'; -import { MockAgentSession } from './mock.js'; -import type { AgentEvent } from './types.js'; +import { MockAgentProtocol } from './mock.js'; +import type { AgentEvent, AgentProtocol } from './types.js'; -describe('MockAgentSession', () => { - it('should yield queued events on send and stream', async () => { - const session = new MockAgentSession(); +const waitForStreamEnd = (session: AgentProtocol): Promise => + new Promise((resolve) => { + const events: AgentEvent[] = []; + const unsubscribe = session.subscribe((e) => { + events.push(e); + if (e.type === 'agent_end') { + unsubscribe(); + resolve(events); + } + }); + }); + +describe('MockAgentProtocol', () => { + it('should emit queued events on send and subscribe', async () => { + const session = new MockAgentProtocol(); const event1 = { type: 'message', role: 'agent', @@ -19,31 +31,30 @@ describe('MockAgentSession', () => { session.pushResponse([event1]); + const streamPromise = waitForStreamEnd(session); + const { streamId } = await session.send({ message: [{ type: 'text', text: 'hi' }], }); expect(streamId).toBeDefined(); - const streamedEvents: AgentEvent[] = []; - for await (const event of session.stream()) { - streamedEvents.push(event); - } + const streamedEvents = await streamPromise; - // Auto stream_start, auto user message, agent message, auto stream_end = 4 events + // Ordered: user message, agent_start, agent message, agent_end = 4 events expect(streamedEvents).toHaveLength(4); - expect(streamedEvents[0].type).toBe('stream_start'); - expect(streamedEvents[1].type).toBe('message'); - expect((streamedEvents[1] as AgentEvent<'message'>).role).toBe('user'); + expect(streamedEvents[0].type).toBe('message'); + expect((streamedEvents[0] as AgentEvent<'message'>).role).toBe('user'); + expect(streamedEvents[1].type).toBe('agent_start'); expect(streamedEvents[2].type).toBe('message'); expect((streamedEvents[2] as AgentEvent<'message'>).role).toBe('agent'); - expect(streamedEvents[3].type).toBe('stream_end'); + expect(streamedEvents[3].type).toBe('agent_end'); expect(session.events).toHaveLength(4); expect(session.events).toEqual(streamedEvents); }); it('should handle multiple responses', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); // Test with empty payload (no message injected) session.pushResponse([]); @@ -57,204 +68,154 @@ describe('MockAgentSession', () => { ]); // First send + const stream1Promise = waitForStreamEnd(session); const { streamId: s1 } = await session.send({ - update: {}, + update: { title: 't1' }, }); - const events1: AgentEvent[] = []; - for await (const e of session.stream()) events1.push(e); - expect(events1).toHaveLength(3); // stream_start, session_update, stream_end - expect(events1[0].type).toBe('stream_start'); - expect(events1[1].type).toBe('session_update'); - expect(events1[2].type).toBe('stream_end'); + const events1 = await stream1Promise; + expect(events1).toHaveLength(3); // session_update, agent_start, agent_end + expect(events1[0].type).toBe('session_update'); + expect(events1[1].type).toBe('agent_start'); + expect(events1[2].type).toBe('agent_end'); // Second send + const stream2Promise = waitForStreamEnd(session); const { streamId: s2 } = await session.send({ - update: {}, + update: { title: 't2' }, }); expect(s1).not.toBe(s2); - const events2: AgentEvent[] = []; - for await (const e of session.stream()) events2.push(e); - expect(events2).toHaveLength(4); // stream_start, session_update, error, stream_end - expect(events2[1].type).toBe('session_update'); + const events2 = await stream2Promise; + expect(events2).toHaveLength(4); // session_update, agent_start, error, agent_end + expect(events2[0].type).toBe('session_update'); + expect(events2[1].type).toBe('agent_start'); expect(events2[2].type).toBe('error'); + expect(events2[3].type).toBe('agent_end'); expect(session.events).toHaveLength(7); }); - it('should allow streaming by streamId', async () => { - const session = new MockAgentSession(); - session.pushResponse([{ type: 'message' }]); - - const { streamId } = await session.send({ - update: {}, - }); + it('should handle abort on a waiting stream', async () => { + const session = new MockAgentProtocol(); + // Use keepOpen to prevent auto agent_end + session.pushResponse([{ type: 'message' }], { keepOpen: true }); const events: AgentEvent[] = []; - for await (const e of session.stream({ streamId })) { + let resolveStream: (evs: AgentEvent[]) => void; + const streamPromise = new Promise((res) => { + resolveStream = res; + }); + + session.subscribe((e) => { events.push(e); - } - expect(events).toHaveLength(4); // start, update, message, end - }); + if (e.type === 'agent_end') { + resolveStream(events); + } + }); - it('should throw when streaming non-existent streamId', async () => { - const session = new MockAgentSession(); - await expect(async () => { - const stream = session.stream({ streamId: 'invalid' }); - await stream.next(); - }).rejects.toThrow('Stream not found: invalid'); - }); + const { streamId: _streamId } = await session.send({ + update: { title: 't' }, + }); - it('should throw when streaming non-existent eventId', async () => { - const session = new MockAgentSession(); - session.pushResponse([{ type: 'message' }]); - await session.send({ update: {} }); - - await expect(async () => { - const stream = session.stream({ eventId: 'invalid' }); - await stream.next(); - }).rejects.toThrow('Event not found: invalid'); - }); - - it('should handle abort on a waiting stream', async () => { - const session = new MockAgentSession(); - // Use keepOpen to prevent auto stream_end - session.pushResponse([{ type: 'message' }], { keepOpen: true }); - const { streamId } = await session.send({ update: {} }); - - const stream = session.stream({ streamId }); - - // Read initial events - const e1 = await stream.next(); - expect(e1.value.type).toBe('stream_start'); - const e2 = await stream.next(); - expect(e2.value.type).toBe('session_update'); - const e3 = await stream.next(); - expect(e3.value.type).toBe('message'); + // Initial events should have been emitted + expect(events.map((e) => e.type)).toEqual([ + 'session_update', + 'agent_start', + 'message', + ]); // At this point, the stream should be "waiting" for more events because it's still active - // and hasn't seen a stream_end. - const abortPromise = session.abort(); - const e4 = await stream.next(); - expect(e4.value.type).toBe('stream_end'); - expect((e4.value as AgentEvent<'stream_end'>).reason).toBe('aborted'); + // and hasn't seen an agent_end. + await session.abort(); - await abortPromise; - expect(await stream.next()).toEqual({ done: true, value: undefined }); + const finalEvents = await streamPromise; + expect(finalEvents[3].type).toBe('agent_end'); + expect((finalEvents[3] as AgentEvent<'agent_end'>).reason).toBe('aborted'); }); it('should handle pushToStream on a waiting stream', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); session.pushResponse([], { keepOpen: true }); - const { streamId } = await session.send({ update: {} }); - const stream = session.stream({ streamId }); - await stream.next(); // start - await stream.next(); // update + const events: AgentEvent[] = []; + session.subscribe((e) => events.push(e)); + + const { streamId } = await session.send({ update: { title: 't' } }); + + expect(events.map((e) => e.type)).toEqual([ + 'session_update', + 'agent_start', + ]); // Push new event to active stream - session.pushToStream(streamId, [{ type: 'message' }]); + session.pushToStream(streamId!, [{ type: 'message' }]); - const e3 = await stream.next(); - expect(e3.value.type).toBe('message'); + expect(events).toHaveLength(3); + expect(events[2].type).toBe('message'); await session.abort(); - const e4 = await stream.next(); - expect(e4.value.type).toBe('stream_end'); + expect(events).toHaveLength(4); + expect(events[3].type).toBe('agent_end'); }); it('should handle pushToStream with close option', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); session.pushResponse([], { keepOpen: true }); - const { streamId } = await session.send({ update: {} }); - const stream = session.stream({ streamId }); - await stream.next(); // start - await stream.next(); // update + const streamPromise = waitForStreamEnd(session); + const { streamId } = await session.send({ update: { title: 't' } }); // Push new event and close - session.pushToStream(streamId, [{ type: 'message' }], { close: true }); + session.pushToStream(streamId!, [{ type: 'message' }], { close: true }); - const e3 = await stream.next(); - expect(e3.value.type).toBe('message'); - - const e4 = await stream.next(); - expect(e4.value.type).toBe('stream_end'); - expect((e4.value as AgentEvent<'stream_end'>).reason).toBe('completed'); - - expect(await stream.next()).toEqual({ done: true, value: undefined }); + const events = await streamPromise; + expect(events.map((e) => e.type)).toEqual([ + 'session_update', + 'agent_start', + 'message', + 'agent_end', + ]); + expect((events[3] as AgentEvent<'agent_end'>).reason).toBe('completed'); }); - it('should not double up on stream_end if provided manually', async () => { - const session = new MockAgentSession(); + it('should not double up on agent_end if provided manually', async () => { + const session = new MockAgentProtocol(); session.pushResponse([ { type: 'message' }, - { type: 'stream_end', reason: 'completed' }, + { type: 'agent_end', reason: 'completed' }, ]); - const { streamId } = await session.send({ update: {} }); - const events: AgentEvent[] = []; - for await (const e of session.stream({ streamId })) { - events.push(e); - } + const streamPromise = waitForStreamEnd(session); + await session.send({ update: { title: 't' } }); - const endEvents = events.filter((e) => e.type === 'stream_end'); + const events = await streamPromise; + const endEvents = events.filter((e) => e.type === 'agent_end'); expect(endEvents).toHaveLength(1); }); - it('should stream after eventId', async () => { - const session = new MockAgentSession(); - // Use manual IDs to test resumption - session.pushResponse([ - { type: 'stream_start', id: 'e1' }, - { type: 'message', id: 'e2' }, - { type: 'stream_end', id: 'e3' }, - ]); - - await session.send({ update: {} }); - - // Stream first event only - const first: AgentEvent[] = []; - for await (const e of session.stream()) { - first.push(e); - if (e.id === 'e1') break; - } - expect(first).toHaveLength(1); - expect(first[0].id).toBe('e1'); - - // Resume from e1 - const second: AgentEvent[] = []; - for await (const e of session.stream({ eventId: 'e1' })) { - second.push(e); - } - expect(second).toHaveLength(3); // update, message, end - expect(second[0].type).toBe('session_update'); - expect(second[1].id).toBe('e2'); - expect(second[2].id).toBe('e3'); - }); - it('should handle elicitations', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); session.pushResponse([]); + const streamPromise = waitForStreamEnd(session); await session.send({ elicitations: [ { requestId: 'r1', action: 'accept', content: { foo: 'bar' } }, ], }); - const events: AgentEvent[] = []; - for await (const e of session.stream()) events.push(e); - - expect(events[1].type).toBe('elicitation_response'); - expect((events[1] as AgentEvent<'elicitation_response'>).requestId).toBe( + const events = await streamPromise; + expect(events[0].type).toBe('elicitation_response'); + expect((events[0] as AgentEvent<'elicitation_response'>).requestId).toBe( 'r1', ); + expect(events[1].type).toBe('agent_start'); }); it('should handle updates and track state', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); session.pushResponse([]); + const streamPromise = waitForStreamEnd(session); await session.send({ update: { title: 'New Title', model: 'gpt-4', config: { x: 1 } }, }); @@ -263,15 +224,24 @@ describe('MockAgentSession', () => { expect(session.model).toBe('gpt-4'); expect(session.config).toEqual({ x: 1 }); - const events: AgentEvent[] = []; - for await (const e of session.stream()) events.push(e); - expect(events[1].type).toBe('session_update'); + const events = await streamPromise; + expect(events[0].type).toBe('session_update'); + expect(events[1].type).toBe('agent_start'); + }); + + it('should return streamId: null if no response queued', async () => { + const session = new MockAgentProtocol(); + const { streamId } = await session.send({ update: { title: 'foo' } }); + expect(streamId).toBeNull(); + expect(session.events).toHaveLength(1); + expect(session.events[0].type).toBe('session_update'); + expect(session.events[0].streamId).toBeNull(); }); it('should throw on action', async () => { - const session = new MockAgentSession(); + const session = new MockAgentProtocol(); await expect( session.send({ action: { type: 'foo', data: {} } }), - ).rejects.toThrow('Actions not supported in MockAgentSession: foo'); + ).rejects.toThrow('Actions not supported in MockAgentProtocol: foo'); }); }); diff --git a/packages/core/src/agent/mock.ts b/packages/core/src/agent/mock.ts index 7baeb61a83..f29e87f878 100644 --- a/packages/core/src/agent/mock.ts +++ b/packages/core/src/agent/mock.ts @@ -9,31 +9,32 @@ import type { AgentEventCommon, AgentEventData, AgentSend, - AgentSession, + AgentProtocol, + Unsubscribe, } from './types.js'; export type MockAgentEvent = Partial & AgentEventData; export interface PushResponseOptions { - /** If true, does not automatically add a stream_end event. */ + /** If true, does not automatically add an agent_end event. */ keepOpen?: boolean; } /** - * A mock implementation of AgentSession for testing. + * A mock implementation of AgentProtocol for testing. * Allows queuing responses that will be yielded when send() is called. */ -export class MockAgentSession implements AgentSession { +export class MockAgentProtocol implements AgentProtocol { private _events: AgentEvent[] = []; private _responses: Array<{ events: MockAgentEvent[]; options?: PushResponseOptions; }> = []; - private _streams = new Map(); + private _subscribers = new Set<(event: AgentEvent) => void>(); private _activeStreamIds = new Set(); - private _lastStreamId?: string; + private _lastStreamId?: string | null; private _nextEventId = 1; - private _streamResolvers = new Map void>>(); + private _nextStreamId = 1; title?: string; model?: string; @@ -50,12 +51,28 @@ export class MockAgentSession implements AgentSession { return this._events; } + subscribe(callback: (event: AgentEvent) => void): Unsubscribe { + this._subscribers.add(callback); + return () => this._subscribers.delete(callback); + } + + private _emit(event: AgentEvent) { + if (!this._events.some((e) => e.id === event.id)) { + this._events.push(event); + } + for (const callback of this._subscribers) { + callback(event); + } + if (event.type === 'agent_end' && event.streamId) { + this._activeStreamIds.delete(event.streamId); + } + } + /** * Queues a sequence of events to be "emitted" by the agent in response to the * next send() call. */ pushResponse(events: MockAgentEvent[], options?: PushResponseOptions) { - // We store them as data and normalize them when send() is called this._responses.push({ events, options }); } @@ -67,11 +84,6 @@ export class MockAgentSession implements AgentSession { events: MockAgentEvent[], options?: { close?: boolean }, ) { - const stream = this._streams.get(streamId); - if (!stream) { - throw new Error(`Stream not found: ${streamId}`); - } - const now = new Date().toISOString(); for (const eventData of events) { const event: AgentEvent = { @@ -80,205 +92,147 @@ export class MockAgentSession implements AgentSession { timestamp: eventData.timestamp ?? now, streamId: eventData.streamId ?? streamId, } as AgentEvent; - stream.push(event); + this._emit(event); } if ( options?.close && - !events.some((eventData) => eventData.type === 'stream_end') + !events.some((eventData) => eventData.type === 'agent_end') ) { - stream.push({ + this._emit({ id: `e-${this._nextEventId++}`, timestamp: now, streamId, - type: 'stream_end', + type: 'agent_end', reason: 'completed', } as AgentEvent); } - - this._notify(streamId); } - private _notify(streamId: string) { - const resolvers = this._streamResolvers.get(streamId); - if (resolvers) { - this._streamResolvers.delete(streamId); - for (const resolve of resolvers) resolve(); - } - } - - async send(payload: AgentSend): Promise<{ streamId: string }> { - const { events: response, options } = this._responses.shift() ?? { + async send(payload: AgentSend): Promise<{ streamId: string | null }> { + const responseData = this._responses.shift(); + const { events: response, options } = responseData ?? { events: [], }; - const streamId = - response[0]?.streamId ?? `mock-stream-${this._streams.size + 1}`; + + // If there were queued responses (even if empty array), we trigger a stream. + const hasResponseEvents = responseData !== undefined; + const streamId = hasResponseEvents + ? (response[0]?.streamId ?? `mock-stream-${this._nextStreamId++}`) + : null; const now = new Date().toISOString(); + const eventsToEmit: AgentEvent[] = []; - if (!response.some((eventData) => eventData.type === 'stream_start')) { - response.unshift({ - type: 'stream_start', - streamId, - }); - } - - const startIndex = response.findIndex( - (eventData) => eventData.type === 'stream_start', - ); + // Helper to normalize and prepare for emission + const normalize = (eventData: MockAgentEvent): AgentEvent => + ({ + ...eventData, + id: eventData.id ?? `e-${this._nextEventId++}`, + timestamp: eventData.timestamp ?? now, + streamId: eventData.streamId ?? streamId, + }) as AgentEvent; + // 1. User/Update event (BEFORE agent_start) if ('message' in payload && payload.message) { - response.splice(startIndex + 1, 0, { - type: 'message', - role: 'user', - content: payload.message, - _meta: payload._meta, - }); - } else if ('elicitations' in payload && payload.elicitations) { - payload.elicitations.forEach((elicitation, i) => { - response.splice(startIndex + 1 + i, 0, { - type: 'elicitation_response', - ...elicitation, + eventsToEmit.push( + normalize({ + type: 'message', + role: 'user', + content: payload.message, _meta: payload._meta, - }); + }), + ); + } else if ('elicitations' in payload && payload.elicitations) { + payload.elicitations.forEach((elicitation) => { + eventsToEmit.push( + normalize({ + type: 'elicitation_response', + ...elicitation, + _meta: payload._meta, + }), + ); }); - } else if ('update' in payload && payload.update) { + } else if ( + 'update' in payload && + payload.update && + Object.keys(payload.update).length > 0 + ) { if (payload.update.title) this.title = payload.update.title; if (payload.update.model) this.model = payload.update.model; if (payload.update.config) { this.config = payload.update.config; } - response.splice(startIndex + 1, 0, { - type: 'session_update', - ...payload.update, - _meta: payload._meta, - }); + eventsToEmit.push( + normalize({ + type: 'session_update', + ...payload.update, + _meta: payload._meta, + }), + ); } else if ('action' in payload && payload.action) { throw new Error( - `Actions not supported in MockAgentSession: ${payload.action.type}`, + `Actions not supported in MockAgentProtocol: ${payload.action.type}`, ); } - if ( - !options?.keepOpen && - !response.some((eventData) => eventData.type === 'stream_end') - ) { - response.push({ - type: 'stream_end', - reason: 'completed', - streamId, - }); - } - - const normalizedResponse: AgentEvent[] = []; - for (const eventData of response) { - const event: AgentEvent = { - ...eventData, - id: eventData.id ?? `e-${this._nextEventId++}`, - timestamp: eventData.timestamp ?? now, - streamId: eventData.streamId ?? streamId, - } as AgentEvent; - normalizedResponse.push(event); - } - - this._streams.set(streamId, normalizedResponse); - this._activeStreamIds.add(streamId); - this._lastStreamId = streamId; - - return { streamId }; - } - - async *stream(options?: { - streamId?: string; - eventId?: string; - }): AsyncIterableIterator { - let streamId = options?.streamId; - - if (options?.eventId) { - const event = this._events.find( - (eventData) => eventData.id === options.eventId, - ); - if (!event) { - throw new Error(`Event not found: ${options.eventId}`); - } - streamId = streamId ?? event.streamId; - } - - streamId = streamId ?? this._lastStreamId; - - if (!streamId) { - return; - } - - const events = this._streams.get(streamId); - if (!events) { - throw new Error(`Stream not found: ${streamId}`); - } - - let i = 0; - if (options?.eventId) { - const idx = events.findIndex( - (eventData) => eventData.id === options.eventId, - ); - if (idx !== -1) { - i = idx + 1; - } else { - // This should theoretically not happen if the event was found in this._events - // but the trajectories match. - throw new Error( - `Event ${options.eventId} not found in stream ${streamId}`, + // 2. agent_start (if stream) + if (streamId) { + if (!response.some((eventData) => eventData.type === 'agent_start')) { + eventsToEmit.push( + normalize({ + type: 'agent_start', + streamId, + }), ); } } - while (true) { - if (i < events.length) { - const event = events[i++]; - // Add to session trajectory if not already present - if (!this._events.some((eventData) => eventData.id === event.id)) { - this._events.push(event); - } - yield event; + // 3. Response events + for (const eventData of response) { + eventsToEmit.push(normalize(eventData)); + } - // If it's a stream_end, we're done with this stream - if (event.type === 'stream_end') { - this._activeStreamIds.delete(streamId); - return; - } - } else { - // No more events in the array currently. Check if we're still active. - if (!this._activeStreamIds.has(streamId)) { - // If we weren't terminated by a stream_end but we're no longer active, - // it was an abort. - const abortEvent: AgentEvent = { - id: `e-${this._nextEventId++}`, - timestamp: new Date().toISOString(), + // 4. agent_end (if stream and not manual) + if (streamId && !options?.keepOpen) { + if (!eventsToEmit.some((e) => e.type === 'agent_end')) { + eventsToEmit.push( + normalize({ + type: 'agent_end', + reason: 'completed', streamId, - type: 'stream_end', - reason: 'aborted', - } as AgentEvent; - if (!this._events.some((e) => e.id === abortEvent.id)) { - this._events.push(abortEvent); - } - yield abortEvent; - return; - } - - // Wait for notification (new event or abort) - await new Promise((resolve) => { - const resolvers = this._streamResolvers.get(streamId) ?? []; - resolvers.push(resolve); - this._streamResolvers.set(streamId, resolvers); - }); + }), + ); } } + + if (streamId) { + this._activeStreamIds.add(streamId); + } + this._lastStreamId = streamId; + + // Emit events asynchronously so the caller receives the streamId first. + if (eventsToEmit.length > 0) { + void Promise.resolve().then(() => { + for (const event of eventsToEmit) { + this._emit(event); + } + }); + } + + return { streamId }; } async abort(): Promise { - if (this._lastStreamId) { + if (this._lastStreamId && this._activeStreamIds.has(this._lastStreamId)) { const streamId = this._lastStreamId; - this._activeStreamIds.delete(streamId); - this._notify(streamId); + this._emit({ + id: `e-${this._nextEventId++}`, + timestamp: new Date().toISOString(), + streamId, + type: 'agent_end', + reason: 'aborted', + } as AgentEvent); } } } diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index 8b698a8e48..3b1c740ad4 100644 --- a/packages/core/src/agent/types.ts +++ b/packages/core/src/agent/types.ts @@ -6,25 +6,27 @@ export type WithMeta = { _meta?: Record }; -export interface AgentSession extends Trajectory { +export type Unsubscribe = () => void; + +export interface AgentProtocol extends Trajectory { /** * Send data to the agent. Promise resolves when action is acknowledged. - * Returns the `streamId` of the stream the message was correlated to -- this may - * be a new stream if idle or an existing stream. - */ - send(payload: AgentSend): Promise<{ streamId: string }>; - /** - * Begin listening to actively streaming data. Stream must have the following - * properties: + * Returns the `streamId` of the stream the message was correlated to -- + * this may be a new stream if idle, an existing stream, or null if no + * stream was triggered. * - * - If no arguments are provided, streams events from an active stream. - * - If a {streamId} is provided, streams ALL events from that stream. - * - If an {eventId} is provided, streams all events AFTER that event. + * When a new stream is created by a send, the streamId MUST be returned + * before the `agent_start` event is emitted for the stream. */ - stream(options?: { - streamId?: string; - eventId?: string; - }): AsyncIterableIterator; + send(payload: AgentSend): Promise<{ streamId: string | null }>; + + /** + * Subscribes the provided callback to all future events emitted by this + * session. Returns an unsubscribe function. + * + * @param callback The callback function to listen to events. + */ + subscribe(callback: (event: AgentEvent) => void): Unsubscribe; /** * Aborts an active stream of agent activity. @@ -32,7 +34,7 @@ export interface AgentSession extends Trajectory { abort(): Promise; /** - * AgentSession implements the Trajectory interface and can retrieve existing events. + * AgentProtocol implements the Trajectory interface and can retrieve existing events. */ readonly events: AgentEvent[]; } @@ -61,7 +63,7 @@ export interface AgentEventCommon { /** Identifies the subagent thread, omitted for "main thread" events. */ threadId?: string; /** Identifies a particular stream of a particular thread. */ - streamId?: string; + streamId?: string | null; /** ISO Timestamp for the time at which the event occurred. */ timestamp: string; /** The concrete type of the event. */ @@ -90,10 +92,10 @@ export interface AgentEvents { session_update: SessionUpdate; /** Message content provided by user, agent, or developer. */ message: Message; - /** Event indicating the start of a new stream. */ - stream_start: StreamStart; - /** Event indicating the end of a running stream. */ - stream_end: StreamEnd; + /** Event indicating the start of agent activity on a stream. */ + agent_start: AgentStart; + /** Event indicating the end of agent activity on a stream. */ + agent_end: AgentEnd; /** Tool request issued by the agent. */ tool_request: ToolRequest; /** Tool update issued by the agent. */ @@ -257,7 +259,7 @@ export interface Usage { cost?: { amount: number; currency?: string }; } -export interface StreamStart { +export interface AgentStart { streamId: string; } @@ -272,7 +274,7 @@ type StreamEndReason = | 'elicitation' | (string & {}); -export interface StreamEnd { +export interface AgentEnd { streamId: string; reason: StreamEndReason; elicitationIds?: string[]; From 5a3c7154df30546dabf96330946e9139c885d13a Mon Sep 17 00:00:00 2001 From: Jerop Kipruto Date: Fri, 20 Mar 2026 10:10:51 -0400 Subject: [PATCH 02/32] feat(worktree): add Git worktree support for isolated parallel sessions (#22973) --- docs/cli/cli-reference.md | 1 + docs/cli/git-worktrees.md | 107 ++++++ docs/cli/session-management.md | 6 + docs/cli/settings.md | 1 + docs/reference/configuration.md | 5 + docs/sidebar.json | 5 + packages/cli/src/config/config.test.ts | 45 +++ packages/cli/src/config/config.ts | 106 +++++- packages/cli/src/config/settings.ts | 4 + packages/cli/src/config/settingsSchema.ts | 10 + packages/cli/src/gemini.test.tsx | 2 + packages/cli/src/gemini.tsx | 10 + packages/cli/src/gemini_cleanup.test.tsx | 2 + .../components/SessionSummaryDisplay.test.tsx | 47 ++- .../ui/components/SessionSummaryDisplay.tsx | 14 +- packages/cli/src/utils/worktreeSetup.test.ts | 124 +++++++ packages/cli/src/utils/worktreeSetup.ts | 43 +++ packages/core/src/config/config.ts | 13 + packages/core/src/index.ts | 1 + .../core/src/services/worktreeService.test.ts | 311 ++++++++++++++++++ packages/core/src/services/worktreeService.ts | 225 +++++++++++++ .../core/src/utils/memoryImportProcessor.ts | 10 +- schemas/settings.schema.json | 7 + 23 files changed, 1090 insertions(+), 9 deletions(-) create mode 100644 docs/cli/git-worktrees.md create mode 100644 packages/cli/src/utils/worktreeSetup.test.ts create mode 100644 packages/cli/src/utils/worktreeSetup.ts create mode 100644 packages/core/src/services/worktreeService.test.ts create mode 100644 packages/core/src/services/worktreeService.ts diff --git a/docs/cli/cli-reference.md b/docs/cli/cli-reference.md index 167801ca05..bc8f8b44ce 100644 --- a/docs/cli/cli-reference.md +++ b/docs/cli/cli-reference.md @@ -50,6 +50,7 @@ These commands are available within the interactive REPL. | `--model` | `-m` | string | `auto` | Model to use. See [Model Selection](#model-selection) for available values. | | `--prompt` | `-p` | string | - | Prompt text. Appended to stdin input if provided. Forces non-interactive mode. | | `--prompt-interactive` | `-i` | string | - | Execute prompt and continue in interactive mode | +| `--worktree` | `-w` | string | - | Start Gemini in a new git worktree. If no name is provided, one is generated automatically. Requires `experimental.worktrees: true` in settings. | | `--sandbox` | `-s` | boolean | `false` | Run in a sandboxed environment for safer execution | | `--approval-mode` | - | string | `default` | Approval mode for tool execution. Choices: `default`, `auto_edit`, `yolo` | | `--yolo` | `-y` | boolean | `false` | **Deprecated.** Auto-approve all actions. Use `--approval-mode=yolo` instead. | diff --git a/docs/cli/git-worktrees.md b/docs/cli/git-worktrees.md new file mode 100644 index 0000000000..5020b3fa9a --- /dev/null +++ b/docs/cli/git-worktrees.md @@ -0,0 +1,107 @@ +# Git Worktrees (experimental) + +When working on multiple tasks at once, you can use Git worktrees to give each +Gemini session its own copy of the codebase. Git worktrees create separate +working directories that each have their own files and branch while sharing the +same repository history. This prevents changes in one session from colliding +with another. + +Learn more about [session management](./session-management.md). + + +> [!NOTE] +> This is an experimental feature currently under active development. Your +> feedback is invaluable as we refine this feature. If you have ideas, +> suggestions, or encounter issues: +> +> - [Open an issue](https://github.com/google-gemini/gemini-cli/issues/new?template=bug_report.yml) on GitHub. +> - Use the **/bug** command within Gemini CLI to file an issue. + +Learn more in the official Git worktree +[documentation](https://git-scm.com/docs/git-worktree). + +## How to enable Git worktrees + +Git worktrees are an experimental feature. You must enable them in your settings +using the `/settings` command or by manually editing your `settings.json` file. + +1. Use the `/settings` command. +2. Search for and set **Enable Git Worktrees** to `true`. + +Alternatively, add the following to your `settings.json`: + +```json +{ + "experimental": { + "worktrees": true + } +} +``` + +## How to use Git worktrees + +Use the `--worktree` (`-w`) flag to create an isolated worktree and start Gemini +CLI in it. + +- **Start with a specific name:** The value you pass becomes both the directory + name (within `.gemini/worktrees/`) and the branch name. + + ```bash + gemini --worktree feature-search + ``` + +- **Start with a random name:** If you omit the name, Gemini generates a random + one automatically (for example, `worktree-a1b2c3d4`). + + ```bash + gemini --worktree + ``` + + +> [!NOTE] +> Remember to initialize your development environment in each new +> worktree according to your project's setup. Depending on your stack, this +> might include running dependency installation (`npm install`, `yarn`), setting +> up virtual environments, or following your project's standard build process. + +## How to exit a Git worktree session + +When you exit a worktree session (using `/quit` or `Ctrl+C`), Gemini leaves the +worktree intact so your work is not lost. This includes your uncommitted changes +(modified files, staged changes, or untracked files) and any new commits you +have made. + +Gemini prioritizes a fast and safe exit: it **does not automatically delete** +your worktree or branch. You are responsible for cleaning up your worktrees +manually once you are finished with them. + +When you exit, Gemini displays instructions on how to resume your work or how to +manually remove the worktree if you no longer need it. + +## Resuming work in a Git worktree + +To resume a session in a worktree, navigate to the worktree directory and start +Gemini CLI with the `--resume` flag and the session ID: + +```bash +cd .gemini/worktrees/feature-search +gemini --resume +``` + +## Managing Git worktrees manually + +For more control over worktree location and branch configuration, or to clean up +a preserved worktree, you can use Git directly: + +- **Clean up a preserved Git worktree:** + ```bash + git worktree remove .gemini/worktrees/feature-search --force + git branch -D worktree-feature-search + ``` +- **Create a Git worktree manually:** + ```bash + git worktree add ../project-feature-search -b feature-search + cd ../project-feature-search && gemini + ``` + +[Open an issue]: https://github.com/google-gemini/gemini-cli/issues diff --git a/docs/cli/session-management.md b/docs/cli/session-management.md index 8e60f61630..74bc4a4337 100644 --- a/docs/cli/session-management.md +++ b/docs/cli/session-management.md @@ -96,6 +96,12 @@ Compatibility aliases: - `/chat ...` works for the same commands. - `/resume checkpoints ...` also remains supported during migration. +## Parallel sessions with Git worktrees + +When working on multiple tasks at once, you can use +[Git worktrees](./git-worktrees.md) to give each Gemini session its own copy of +the codebase. This prevents changes in one session from colliding with another. + ## Managing sessions You can list and delete sessions to keep your history organized and manage disk diff --git a/docs/cli/settings.md b/docs/cli/settings.md index 85373f1034..ead0050fbd 100644 --- a/docs/cli/settings.md +++ b/docs/cli/settings.md @@ -151,6 +151,7 @@ they appear in the UI. | UI Label | Setting | Description | Default | | -------------------------- | ---------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | Enable Tool Output Masking | `experimental.toolOutputMasking.enabled` | Enables tool output masking to save tokens. | `true` | +| Enable Git Worktrees | `experimental.worktrees` | Enable automated Git worktree management for parallel work. | `false` | | Use OSC 52 Paste | `experimental.useOSC52Paste` | Use OSC 52 for pasting. This may be more robust than the default system when using remote terminal sessions (if your terminal is configured to allow it). | `false` | | Use OSC 52 Copy | `experimental.useOSC52Copy` | Use OSC 52 for copying. This may be more robust than the default system when using remote terminal sessions (if your terminal is configured to allow it). | `false` | | Plan | `experimental.plan` | Enable Plan Mode. | `true` | diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index d3b08d565a..5791bbf457 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -1527,6 +1527,11 @@ their corresponding top-level category object in your `settings.json` file. - **Default:** `true` - **Requires restart:** Yes +- **`experimental.worktrees`** (boolean): + - **Description:** Enable automated Git worktree management for parallel work. + - **Default:** `false` + - **Requires restart:** Yes + - **`experimental.extensionManagement`** (boolean): - **Description:** Enable extension management features. - **Default:** `true` diff --git a/docs/sidebar.json b/docs/sidebar.json index 6cac5ec9fd..7198a0336b 100644 --- a/docs/sidebar.json +++ b/docs/sidebar.json @@ -99,6 +99,11 @@ { "label": "Agent Skills", "slug": "docs/cli/skills" }, { "label": "Checkpointing", "slug": "docs/cli/checkpointing" }, { "label": "Headless mode", "slug": "docs/cli/headless" }, + { + "label": "Git worktrees", + "badge": "🔬", + "slug": "docs/cli/git-worktrees" + }, { "label": "Hooks", "collapsed": true, diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index c046f0c0e7..746fc14475 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -226,6 +226,51 @@ afterEach(() => { }); describe('parseArguments', () => { + describe('worktree', () => { + it('should parse --worktree flag when provided with a name', async () => { + process.argv = ['node', 'script.js', '--worktree', 'my-feature']; + const settings = createTestMergedSettings(); + settings.experimental.worktrees = true; + const argv = await parseArguments(settings); + expect(argv.worktree).toBe('my-feature'); + }); + + it('should generate a random name when --worktree is provided without a name', async () => { + process.argv = ['node', 'script.js', '--worktree']; + const settings = createTestMergedSettings(); + settings.experimental.worktrees = true; + const argv = await parseArguments(settings); + expect(argv.worktree).toBeDefined(); + expect(argv.worktree).not.toBe(''); + expect(typeof argv.worktree).toBe('string'); + }); + + it('should throw an error when --worktree is used but experimental.worktrees is not enabled', async () => { + process.argv = ['node', 'script.js', '--worktree', 'feature']; + const settings = createTestMergedSettings(); + settings.experimental.worktrees = false; + + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit called'); + }); + const mockConsoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + + await expect(parseArguments(settings)).rejects.toThrow( + 'process.exit called', + ); + expect(mockConsoleError).toHaveBeenCalledWith( + expect.stringContaining( + 'The --worktree flag is only available when experimental.worktrees is enabled in your settings.', + ), + ); + + mockExit.mockRestore(); + mockConsoleError.mockRestore(); + }); + }); + it.each([ { description: 'long flags', diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index fdcd18c086..227ad4e8ed 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -4,10 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import yargs from 'yargs/yargs'; +import yargs from 'yargs'; import { hideBin } from 'yargs/helpers'; import process from 'node:process'; import * as path from 'node:path'; +import { execa } from 'execa'; import { mcpCommand } from '../commands/mcp.js'; import { extensionsCommand } from '../commands/extensions.js'; import { skillsCommand } from '../commands/skills.js'; @@ -38,6 +39,9 @@ import { applyAdminAllowlist, applyRequiredServers, getAdminBlockedMcpServersMessage, + getProjectRootForWorktree, + isGeminiWorktree, + type WorktreeSettings, type HookDefinition, type HookEventName, type OutputFormat, @@ -48,6 +52,8 @@ import { type MergedSettings, saveModelChange, loadSettings, + isWorktreeEnabled, + type LoadedSettings, } from './settings.js'; import { loadSandboxConfig } from './sandboxConfig.js'; @@ -74,6 +80,7 @@ export interface CliArgs { debug: boolean | undefined; prompt: string | undefined; promptInteractive: string | undefined; + worktree?: string; yolo: boolean | undefined; approvalMode: string | undefined; @@ -115,6 +122,36 @@ const coerceCommaSeparated = (values: string[]): string[] => { ); }; +/** + * Pre-parses the command line arguments to find the worktree flag. + * Used for early setup before full argument parsing with settings. + */ +export function getWorktreeArg(argv: string[]): string | undefined { + const result = yargs(hideBin(argv)) + .help(false) + .version(false) + .option('worktree', { alias: 'w', type: 'string' }) + .strict(false) + .exitProcess(false) + .parseSync(); + + if (result.worktree === undefined) return undefined; + return typeof result.worktree === 'string' ? result.worktree.trim() : ''; +} + +/** + * Checks if a worktree is requested via CLI and enabled in settings. + * Returns the requested name (can be empty string for auto-generated) or undefined. + */ +export function getRequestedWorktreeName( + settings: LoadedSettings, +): string | undefined { + if (!isWorktreeEnabled(settings)) { + return undefined; + } + return getWorktreeArg(process.argv); +} + export async function parseArguments( settings: MergedSettings, ): Promise { @@ -158,6 +195,20 @@ export async function parseArguments( description: 'Execute the provided prompt and continue in interactive mode', }) + .option('worktree', { + alias: 'w', + type: 'string', + skipValidation: true, + description: + 'Start Gemini in a new git worktree. If no name is provided, one is generated automatically.', + coerce: (value: unknown): string => { + const trimmed = typeof value === 'string' ? value.trim() : ''; + if (trimmed === '') { + return Math.random().toString(36).substring(2, 10); + } + return trimmed; + }, + }) .option('sandbox', { alias: 's', type: 'boolean', @@ -335,6 +386,9 @@ export async function parseArguments( ) { return `Invalid values:\n Argument: output-format, Given: "${argv['outputFormat']}", Choices: "text", "json", "stream-json"`; } + if (argv['worktree'] && !settings.experimental?.worktrees) { + return 'The --worktree flag is only available when experimental.worktrees is enabled in your settings.'; + } return true; }); @@ -420,6 +474,7 @@ export interface LoadCliConfigOptions { projectHooks?: { [K in HookEventName]?: HookDefinition[] } & { disabled?: string[]; }; + worktreeSettings?: WorktreeSettings; } export async function loadCliConfig( @@ -431,6 +486,9 @@ export async function loadCliConfig( const { cwd = process.cwd(), projectHooks } = options; const debugMode = isDebugMode(argv); + const worktreeSettings = + options.worktreeSettings ?? (await resolveWorktreeSettings(cwd)); + if (argv.sandbox) { process.env['GEMINI_SANDBOX'] = 'true'; } @@ -802,6 +860,7 @@ export async function loadCliConfig( importFormat: settings.context?.importFormat, debugMode, question, + worktreeSettings, coreTools: settings.tools?.core || undefined, allowedTools: allowedTools.length > 0 ? allowedTools : undefined, @@ -943,3 +1002,48 @@ function mergeExcludeTools( ]); return Array.from(allExcludeTools); } + +async function resolveWorktreeSettings( + cwd: string, +): Promise { + let worktreePath: string | undefined; + try { + const { stdout } = await execa('git', ['rev-parse', '--show-toplevel'], { + cwd, + }); + const toplevel = stdout.trim(); + const projectRoot = await getProjectRootForWorktree(toplevel); + + if (isGeminiWorktree(toplevel, projectRoot)) { + worktreePath = toplevel; + } + } catch (_e) { + return undefined; + } + + if (!worktreePath) { + return undefined; + } + + let worktreeBaseSha: string | undefined; + try { + const { stdout } = await execa('git', ['rev-parse', 'HEAD'], { + cwd: worktreePath, + }); + worktreeBaseSha = stdout.trim(); + } catch (e: unknown) { + debugLogger.debug( + `Failed to resolve worktree base SHA at ${worktreePath}: ${e instanceof Error ? e.message : String(e)}`, + ); + } + + if (!worktreeBaseSha) { + return undefined; + } + + return { + name: path.basename(worktreePath), + path: worktreePath, + baseSha: worktreeBaseSha, + }; +} diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index beecd6a017..984bdb8d60 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -632,6 +632,10 @@ export function resetSettingsCacheForTesting() { settingsCache.clear(); } +export function isWorktreeEnabled(settings: LoadedSettings): boolean { + return settings.merged.experimental.worktrees; +} + /** * Loads settings from user and workspace directories. * Project settings override user settings. diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index f1711f3b92..3724253e97 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1906,6 +1906,16 @@ const SETTINGS_SCHEMA = { description: 'Enable local and remote subagents.', showInDialog: false, }, + worktrees: { + type: 'boolean', + label: 'Enable Git Worktrees', + category: 'Experimental', + requiresRestart: true, + default: false, + description: + 'Enable automated Git worktree management for parallel work.', + showInDialog: true, + }, extensionManagement: { type: 'boolean', label: 'Extension Management', diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 31fec36db0..08c2cbabe8 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -199,6 +199,8 @@ vi.mock('./config/config.js', () => ({ networkAccess: false, }), isDebugMode: vi.fn(() => false), + getRequestedWorktreeName: vi.fn(() => undefined), + getWorktreeArg: vi.fn(() => undefined), })); vi.mock('read-package-up', () => ({ diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 4722bb73f3..c8cd2b3cd8 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -9,6 +9,7 @@ import { WarningPriority, type Config, type ResumedSessionData, + type WorktreeInfo, type OutputPayload, type ConsoleLogPayload, type UserFeedbackPayload, @@ -63,6 +64,7 @@ import { registerTelemetryConfig, setupSignalHandlers, } from './utils/cleanup.js'; +import { setupWorktree } from './utils/worktreeSetup.js'; import { cleanupToolOutputFiles, cleanupExpiredSessions, @@ -210,6 +212,13 @@ export async function main() { const settings = loadSettings(); loadSettingsHandle?.end(); + // If a worktree is requested and enabled, set it up early. + const requestedWorktree = cliConfig.getRequestedWorktreeName(settings); + let worktreeInfo: WorktreeInfo | undefined; + if (requestedWorktree !== undefined) { + worktreeInfo = await setupWorktree(requestedWorktree || undefined); + } + // Report settings errors once during startup settings.errors.forEach((error) => { coreEvents.emitFeedback('warning', error.message); @@ -426,6 +435,7 @@ export async function main() { const loadConfigHandle = startupProfiler.start('load_cli_config'); const config = await loadCliConfig(settings.merged, sessionId, argv, { projectHooks: settings.workspace.settings.hooks, + worktreeSettings: worktreeInfo, }); loadConfigHandle?.end(); diff --git a/packages/cli/src/gemini_cleanup.test.tsx b/packages/cli/src/gemini_cleanup.test.tsx index 9be9fc6194..382ad3f81f 100644 --- a/packages/cli/src/gemini_cleanup.test.tsx +++ b/packages/cli/src/gemini_cleanup.test.tsx @@ -72,6 +72,8 @@ vi.mock('./config/config.js', () => ({ } as unknown as Config), parseArguments: vi.fn().mockResolvedValue({}), isDebugMode: vi.fn(() => false), + getRequestedWorktreeName: vi.fn(() => undefined), + getWorktreeArg: vi.fn(() => undefined), })); vi.mock('read-package-up', () => ({ diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx index 9c811fc741..f5d1ebbd5e 100644 --- a/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx +++ b/packages/cli/src/ui/components/SessionSummaryDisplay.test.tsx @@ -8,10 +8,12 @@ import { renderWithProviders } from '../../test-utils/render.js'; import { describe, it, expect, vi, beforeEach } from 'vitest'; import { SessionSummaryDisplay } from './SessionSummaryDisplay.js'; import * as SessionContext from '../contexts/SessionContext.js'; +import { useConfig } from '../contexts/ConfigContext.js'; import { type SessionMetrics } from '../contexts/SessionContext.js'; import { ToolCallDecision, getShellConfiguration, + type WorktreeSettings, } from '@google/gemini-cli-core'; vi.mock('@google/gemini-cli-core', async (importOriginal) => { @@ -24,19 +26,30 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }); vi.mock('../contexts/SessionContext.js', async (importOriginal) => { - const actual = await importOriginal(); + const actual = + await importOriginal(); return { ...actual, useSessionStats: vi.fn(), }; }); +vi.mock('../contexts/ConfigContext.js', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + useConfig: vi.fn(), + }; +}); + const getShellConfigurationMock = vi.mocked(getShellConfiguration); const useSessionStatsMock = vi.mocked(SessionContext.useSessionStats); const renderWithMockedStats = async ( metrics: SessionMetrics, sessionId = 'test-session', + worktreeSettings?: WorktreeSettings, ) => { useSessionStatsMock.mockReturnValue({ stats: { @@ -49,7 +62,11 @@ const renderWithMockedStats = async ( getPromptCount: () => 5, startNewPrompt: vi.fn(), - }); + } as unknown as ReturnType); + + vi.mocked(useConfig).mockReturnValue({ + getWorktreeSettings: () => worktreeSettings, + } as never); const result = await renderWithProviders( , @@ -188,4 +205,30 @@ describe('', () => { unmount(); }); }); + + describe('Worktree status', () => { + it('renders worktree instructions when worktreeSettings are present', async () => { + const worktreeSettings: WorktreeSettings = { + name: 'foo-bar', + path: '/path/to/foo-bar', + baseSha: 'base-sha', + }; + + const { lastFrame, unmount } = await renderWithMockedStats( + emptyMetrics, + 'test-session', + worktreeSettings, + ); + const output = lastFrame(); + + expect(output).toContain('To resume work in this worktree:'); + expect(output).toContain( + 'cd /path/to/foo-bar && gemini --resume test-session', + ); + expect(output).toContain( + 'To remove manually: git worktree remove /path/to/foo-bar', + ); + unmount(); + }); + }); }); diff --git a/packages/cli/src/ui/components/SessionSummaryDisplay.tsx b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx index 5b0a461682..7313949a9c 100644 --- a/packages/cli/src/ui/components/SessionSummaryDisplay.tsx +++ b/packages/cli/src/ui/components/SessionSummaryDisplay.tsx @@ -7,6 +7,7 @@ import type React from 'react'; import { StatsDisplay } from './StatsDisplay.js'; import { useSessionStats } from '../contexts/SessionContext.js'; +import { useConfig } from '../contexts/ConfigContext.js'; import { escapeShellArg, getShellConfiguration } from '@google/gemini-cli-core'; interface SessionSummaryDisplayProps { @@ -17,8 +18,19 @@ export const SessionSummaryDisplay: React.FC = ({ duration, }) => { const { stats } = useSessionStats(); + const config = useConfig(); const { shell } = getShellConfiguration(); - const footer = `To resume this session: gemini --resume ${escapeShellArg(stats.sessionId, shell)}`; + + const worktreeSettings = config.getWorktreeSettings(); + + const escapedSessionId = escapeShellArg(stats.sessionId, shell); + let footer = `To resume this session: gemini --resume ${escapedSessionId}`; + + if (worktreeSettings) { + footer = + `To resume work in this worktree: cd ${escapeShellArg(worktreeSettings.path, shell)} && gemini --resume ${escapedSessionId}\n` + + `To remove manually: git worktree remove ${escapeShellArg(worktreeSettings.path, shell)}`; + } return ( { + const actual = + await importOriginal(); + return { + ...actual, + getProjectRootForWorktree: vi.fn(), + createWorktreeService: vi.fn(), + debugLogger: { + log: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }, + writeToStdout: vi.fn(), + writeToStderr: vi.fn(), + }; +}); + +describe('setupWorktree', () => { + const originalEnv = { ...process.env }; + const originalCwd = process.cwd; + + const mockService = { + setup: vi.fn(), + maybeCleanup: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...originalEnv }; + + // Mock process.cwd and process.chdir + let currentPath = '/mock/project'; + process.cwd = vi.fn().mockImplementation(() => currentPath); + process.chdir = vi.fn().mockImplementation((newPath) => { + currentPath = newPath; + }); + + // Mock successful execution of core utilities + vi.mocked(coreFunctions.getProjectRootForWorktree).mockResolvedValue( + '/mock/project', + ); + vi.mocked(coreFunctions.createWorktreeService).mockResolvedValue( + mockService as never, + ); + mockService.setup.mockResolvedValue({ + name: 'my-feature', + path: '/mock/project/.gemini/worktrees/my-feature', + baseSha: 'base-sha', + }); + }); + + afterEach(() => { + process.env = { ...originalEnv }; + process.cwd = originalCwd; + delete (process as { chdir?: typeof process.chdir }).chdir; + }); + + it('should create and switch to a new worktree', async () => { + await setupWorktree('my-feature'); + + expect(coreFunctions.getProjectRootForWorktree).toHaveBeenCalledWith( + '/mock/project', + ); + expect(coreFunctions.createWorktreeService).toHaveBeenCalledWith( + '/mock/project', + ); + expect(mockService.setup).toHaveBeenCalledWith('my-feature'); + expect(process.chdir).toHaveBeenCalledWith( + '/mock/project/.gemini/worktrees/my-feature', + ); + expect(process.env['GEMINI_CLI_WORKTREE_HANDLED']).toBe('1'); + }); + + it('should generate a name if worktreeName is undefined', async () => { + mockService.setup.mockResolvedValue({ + name: 'generated-name', + path: '/mock/project/.gemini/worktrees/generated-name', + baseSha: 'base-sha', + }); + + await setupWorktree(undefined); + + expect(mockService.setup).toHaveBeenCalledWith(undefined); + }); + + it('should skip worktree creation if GEMINI_CLI_WORKTREE_HANDLED is set', async () => { + process.env['GEMINI_CLI_WORKTREE_HANDLED'] = '1'; + + await setupWorktree('my-feature'); + + expect(coreFunctions.createWorktreeService).not.toHaveBeenCalled(); + expect(process.chdir).not.toHaveBeenCalled(); + }); + + it('should handle errors gracefully and exit', async () => { + const mockExit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('PROCESS_EXIT'); + }); + + mockService.setup.mockRejectedValue(new Error('Git failure')); + + await expect(setupWorktree('my-feature')).rejects.toThrow('PROCESS_EXIT'); + + expect(coreFunctions.writeToStderr).toHaveBeenCalledWith( + expect.stringContaining( + 'Failed to create or switch to worktree: Git failure', + ), + ); + expect(mockExit).toHaveBeenCalledWith(1); + + mockExit.mockRestore(); + }); +}); diff --git a/packages/cli/src/utils/worktreeSetup.ts b/packages/cli/src/utils/worktreeSetup.ts new file mode 100644 index 0000000000..596c367d3e --- /dev/null +++ b/packages/cli/src/utils/worktreeSetup.ts @@ -0,0 +1,43 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + getProjectRootForWorktree, + createWorktreeService, + writeToStderr, + type WorktreeInfo, +} from '@google/gemini-cli-core'; + +/** + * Sets up a git worktree for parallel sessions. + * + * This function uses a guard (GEMINI_CLI_WORKTREE_HANDLED) to ensure that + * when the CLI relaunches itself (e.g. for memory allocation), it doesn't + * attempt to create a nested worktree. + */ +export async function setupWorktree( + worktreeName: string | undefined, +): Promise { + if (process.env['GEMINI_CLI_WORKTREE_HANDLED'] === '1') { + return undefined; + } + + try { + const projectRoot = await getProjectRootForWorktree(process.cwd()); + const service = await createWorktreeService(projectRoot); + + const worktreeInfo = await service.setup(worktreeName || undefined); + + process.chdir(worktreeInfo.path); + process.env['GEMINI_CLI_WORKTREE_HANDLED'] = '1'; + + return worktreeInfo; + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + writeToStderr(`Failed to create or switch to worktree: ${errorMessage}\n`); + process.exit(1); + } +} diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 5bac6d086c..eb2c3f90f1 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -528,6 +528,12 @@ export interface PolicyUpdateConfirmationRequest { newHash: string; } +export interface WorktreeSettings { + name: string; + path: string; + baseSha: string; +} + export interface ConfigParameters { sessionId: string; clientName?: string; @@ -651,6 +657,7 @@ export interface ConfigParameters { plan?: boolean; tracker?: boolean; planSettings?: PlanSettings; + worktreeSettings?: WorktreeSettings; modelSteering?: boolean; onModelChange?: (model: string) => void; mcpEnabled?: boolean; @@ -695,6 +702,7 @@ export class Config implements McpContext, AgentLoopContext { private workspaceContext: WorkspaceContext; private readonly debugMode: boolean; private readonly question: string | undefined; + private readonly worktreeSettings: WorktreeSettings | undefined; readonly enableConseca: boolean; private readonly coreTools: string[] | undefined; @@ -925,6 +933,7 @@ export class Config implements McpContext, AgentLoopContext { this.pendingIncludeDirectories = params.includeDirectories ?? []; this.debugMode = params.debugMode; this.question = params.question; + this.worktreeSettings = params.worktreeSettings; this.coreTools = params.coreTools; this.mainAgentTools = params.mainAgentTools; @@ -1555,6 +1564,10 @@ export class Config implements McpContext, AgentLoopContext { return this.promptId; } + getWorktreeSettings(): WorktreeSettings | undefined { + return this.worktreeSettings; + } + getClientName(): string | undefined { return this.clientName; } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 32572c86a0..5729730365 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -237,6 +237,7 @@ export * from './agents/types.js'; // Export stdio utils export * from './utils/stdio.js'; export * from './utils/terminal.js'; +export * from './services/worktreeService.js'; // Export voice utilities export * from './voice/responseFormatter.js'; diff --git a/packages/core/src/services/worktreeService.test.ts b/packages/core/src/services/worktreeService.test.ts new file mode 100644 index 0000000000..b3d831e6b4 --- /dev/null +++ b/packages/core/src/services/worktreeService.test.ts @@ -0,0 +1,311 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import * as path from 'node:path'; +import * as fs from 'node:fs/promises'; +import { + getProjectRootForWorktree, + createWorktree, + isGeminiWorktree, + hasWorktreeChanges, + cleanupWorktree, + getWorktreePath, + WorktreeService, +} from './worktreeService.js'; +import { execa } from 'execa'; + +vi.mock('execa'); +vi.mock('node:fs/promises'); +vi.mock('node:fs', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + realpathSync: vi.fn((p: string) => p), + }; +}); + +describe('worktree utilities', () => { + const projectRoot = '/mock/project'; + const worktreeName = 'test-feature'; + const expectedPath = path.join( + projectRoot, + '.gemini', + 'worktrees', + worktreeName, + ); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('getProjectRootForWorktree', () => { + it('should return the project root from git common dir', async () => { + // In main repo, git-common-dir is often just ".git" + vi.mocked(execa).mockResolvedValue({ + stdout: '.git\n', + } as never); + + const result = await getProjectRootForWorktree('/mock/project'); + expect(result).toBe('/mock/project'); + expect(execa).toHaveBeenCalledWith( + 'git', + ['rev-parse', '--git-common-dir'], + { cwd: '/mock/project' }, + ); + }); + + it('should resolve absolute git common dir paths (as seen in worktrees)', async () => { + // Inside a worktree, git-common-dir is usually an absolute path to the main .git folder + vi.mocked(execa).mockResolvedValue({ + stdout: '/mock/project/.git\n', + } as never); + + const result = await getProjectRootForWorktree( + '/mock/project/.gemini/worktrees/my-feature', + ); + expect(result).toBe('/mock/project'); + }); + + it('should fallback to cwd if git command fails', async () => { + vi.mocked(execa).mockRejectedValue(new Error('not a git repo')); + + const result = await getProjectRootForWorktree('/mock/non-git/src'); + expect(result).toBe('/mock/non-git/src'); + }); + }); + + describe('getWorktreePath', () => { + it('should return the correct path for a given name', () => { + expect(getWorktreePath(projectRoot, worktreeName)).toBe(expectedPath); + }); + }); + + describe('createWorktree', () => { + it('should execute git worktree add with correct branch and path', async () => { + vi.mocked(execa).mockResolvedValue({ stdout: '' } as never); + + const resultPath = await createWorktree(projectRoot, worktreeName); + + expect(resultPath).toBe(expectedPath); + expect(execa).toHaveBeenCalledWith( + 'git', + ['worktree', 'add', expectedPath, '-b', `worktree-${worktreeName}`], + { cwd: projectRoot }, + ); + }); + + it('should throw an error if git worktree add fails', async () => { + vi.mocked(execa).mockRejectedValue(new Error('git failed')); + + await expect(createWorktree(projectRoot, worktreeName)).rejects.toThrow( + 'git failed', + ); + }); + }); + + describe('isGeminiWorktree', () => { + it('should return true for a valid gemini worktree path', () => { + expect(isGeminiWorktree(expectedPath, projectRoot)).toBe(true); + expect( + isGeminiWorktree(path.join(expectedPath, 'src'), projectRoot), + ).toBe(true); + }); + + it('should return false for a path outside gemini worktrees', () => { + expect(isGeminiWorktree(path.join(projectRoot, 'src'), projectRoot)).toBe( + false, + ); + expect(isGeminiWorktree('/some/other/path', projectRoot)).toBe(false); + }); + }); + + describe('hasWorktreeChanges', () => { + it('should return true if git status --porcelain has output', async () => { + vi.mocked(execa).mockResolvedValue({ + stdout: ' M somefile.txt\n?? newfile.txt', + } as never); + + const hasChanges = await hasWorktreeChanges(expectedPath); + + expect(hasChanges).toBe(true); + expect(execa).toHaveBeenCalledWith('git', ['status', '--porcelain'], { + cwd: expectedPath, + }); + }); + + it('should return true if there are untracked files', async () => { + vi.mocked(execa).mockResolvedValue({ + stdout: '?? untracked-file.txt\n', + } as never); + + const hasChanges = await hasWorktreeChanges(expectedPath); + + expect(hasChanges).toBe(true); + }); + + it('should return true if HEAD differs from baseSha', async () => { + vi.mocked(execa) + .mockResolvedValueOnce({ stdout: '' } as never) // status clean + .mockResolvedValueOnce({ stdout: 'different-sha' } as never); // HEAD moved + + const hasChanges = await hasWorktreeChanges(expectedPath, 'base-sha'); + + expect(hasChanges).toBe(true); + }); + + it('should return false if status is clean and HEAD matches baseSha', async () => { + vi.mocked(execa) + .mockResolvedValueOnce({ stdout: '' } as never) // status clean + .mockResolvedValueOnce({ stdout: 'base-sha' } as never); // HEAD same + + const hasChanges = await hasWorktreeChanges(expectedPath, 'base-sha'); + + expect(hasChanges).toBe(false); + }); + + it('should return true if any git command fails', async () => { + vi.mocked(execa).mockRejectedValue(new Error('git error')); + + const hasChanges = await hasWorktreeChanges(expectedPath); + + expect(hasChanges).toBe(true); + }); + }); + + describe('cleanupWorktree', () => { + it('should remove the worktree and delete the branch', async () => { + vi.mocked(fs.access).mockResolvedValue(undefined); + vi.mocked(execa) + .mockResolvedValueOnce({ + stdout: `worktree-${worktreeName}\n`, + } as never) // branch --show-current + .mockResolvedValueOnce({ stdout: '' } as never) // remove + .mockResolvedValueOnce({ stdout: '' } as never); // branch -D + + await cleanupWorktree(expectedPath, projectRoot); + + expect(execa).toHaveBeenCalledTimes(3); + expect(execa).toHaveBeenNthCalledWith( + 1, + 'git', + ['-C', expectedPath, 'branch', '--show-current'], + { cwd: projectRoot }, + ); + expect(execa).toHaveBeenNthCalledWith( + 2, + 'git', + ['worktree', 'remove', expectedPath, '--force'], + { cwd: projectRoot }, + ); + expect(execa).toHaveBeenNthCalledWith( + 3, + 'git', + ['branch', '-D', `worktree-${worktreeName}`], + { cwd: projectRoot }, + ); + }); + + it('should handle branch discovery failure gracefully', async () => { + vi.mocked(fs.access).mockResolvedValue(undefined); + vi.mocked(execa) + .mockResolvedValueOnce({ stdout: '' } as never) // no branch found + .mockResolvedValueOnce({ stdout: '' } as never); // remove + + await cleanupWorktree(expectedPath, projectRoot); + + expect(execa).toHaveBeenCalledTimes(2); + expect(execa).toHaveBeenNthCalledWith( + 2, + 'git', + ['worktree', 'remove', expectedPath, '--force'], + { cwd: projectRoot }, + ); + }); + }); +}); + +describe('WorktreeService', () => { + const projectRoot = '/mock/project'; + const service = new WorktreeService(projectRoot); + + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('setup', () => { + it('should capture baseSha and create a worktree', async () => { + vi.mocked(execa).mockResolvedValue({ + stdout: 'current-sha\n', + } as never); + + const info = await service.setup('feature-x'); + + expect(execa).toHaveBeenCalledWith('git', ['rev-parse', 'HEAD'], { + cwd: projectRoot, + }); + expect(info.name).toBe('feature-x'); + expect(info.baseSha).toBe('current-sha'); + expect(info.path).toContain('feature-x'); + }); + + it('should generate a timestamped name if none provided', async () => { + vi.mocked(execa).mockResolvedValue({ + stdout: 'current-sha\n', + } as never); + + const info = await service.setup(); + + expect(info.name).toMatch(/^\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}-\w+/); + expect(info.path).toContain(info.name); + }); + }); + + describe('maybeCleanup', () => { + const info = { + name: 'feature-x', + path: '/mock/project/.gemini/worktrees/feature-x', + baseSha: 'base-sha', + }; + + it('should cleanup unmodified worktrees', async () => { + // Mock hasWorktreeChanges -> false (no changes) + vi.mocked(execa) + .mockResolvedValueOnce({ stdout: '' } as never) // status check + .mockResolvedValueOnce({ stdout: 'base-sha' } as never); // SHA check + + vi.mocked(fs.access).mockResolvedValue(undefined); + vi.mocked(execa).mockResolvedValue({ stdout: '' } as never); // cleanup calls + + const cleanedUp = await service.maybeCleanup(info); + + expect(cleanedUp).toBe(true); + // Verify cleanupWorktree utilities were called (execa calls inside cleanupWorktree) + expect(execa).toHaveBeenCalledWith( + expect.anything(), + expect.arrayContaining(['worktree', 'remove', info.path, '--force']), + expect.anything(), + ); + }); + + it('should preserve modified worktrees', async () => { + // Mock hasWorktreeChanges -> true (changes detected) + vi.mocked(execa).mockResolvedValue({ + stdout: ' M modified-file.ts', + } as never); + + const cleanedUp = await service.maybeCleanup(info); + + expect(cleanedUp).toBe(false); + // Ensure cleanupWorktree was NOT called + expect(execa).not.toHaveBeenCalledWith( + expect.anything(), + expect.arrayContaining(['worktree', 'remove']), + expect.anything(), + ); + }); + }); +}); diff --git a/packages/core/src/services/worktreeService.ts b/packages/core/src/services/worktreeService.ts new file mode 100644 index 0000000000..0b6bd20648 --- /dev/null +++ b/packages/core/src/services/worktreeService.ts @@ -0,0 +1,225 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as path from 'node:path'; +import * as fs from 'node:fs/promises'; +import { realpathSync } from 'node:fs'; +import { execa } from 'execa'; +import { debugLogger } from '../utils/debugLogger.js'; + +export interface WorktreeInfo { + name: string; + path: string; + baseSha: string; +} + +/** + * Service for managing Git worktrees within Gemini CLI. + * Handles creation, cleanup, and environment setup for isolated sessions. + */ +export class WorktreeService { + constructor(private readonly projectRoot: string) {} + + /** + * Creates a new worktree and prepares the environment. + */ + async setup(name?: string): Promise { + let worktreeName = name?.trim(); + + if (!worktreeName) { + const now = new Date(); + const timestamp = now + .toISOString() + .replace(/[:.]/g, '-') + .replace('T', '-') + .replace('Z', ''); + const randomSuffix = Math.random().toString(36).substring(2, 6); + worktreeName = `${timestamp}-${randomSuffix}`; + } + + // Capture the base commit before creating the worktree + const { stdout: baseSha } = await execa('git', ['rev-parse', 'HEAD'], { + cwd: this.projectRoot, + }); + + const worktreePath = await createWorktree(this.projectRoot, worktreeName); + + return { + name: worktreeName, + path: worktreePath, + baseSha: baseSha.trim(), + }; + } + + /** + * Checks if a worktree has changes and cleans it up if it's unmodified. + */ + async maybeCleanup(info: WorktreeInfo): Promise { + const hasChanges = await hasWorktreeChanges(info.path, info.baseSha); + + if (!hasChanges) { + try { + await cleanupWorktree(info.path, this.projectRoot); + debugLogger.log( + `Automatically cleaned up unmodified worktree: ${info.path}`, + ); + return true; + } catch (error) { + debugLogger.error( + `Failed to clean up worktree ${info.path}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } else { + debugLogger.debug( + `Preserving worktree ${info.path} because it has changes.`, + ); + } + + return false; + } +} + +export async function createWorktreeService( + cwd: string, +): Promise { + const projectRoot = await getProjectRootForWorktree(cwd); + return new WorktreeService(projectRoot); +} + +// Low-level worktree utilities + +export async function getProjectRootForWorktree(cwd: string): Promise { + try { + const { stdout } = await execa('git', ['rev-parse', '--git-common-dir'], { + cwd, + }); + const gitCommonDir = stdout.trim(); + const absoluteGitDir = path.isAbsolute(gitCommonDir) + ? gitCommonDir + : path.resolve(cwd, gitCommonDir); + + // The project root is the parent of the .git directory/file + return path.dirname(absoluteGitDir); + } catch (e: unknown) { + debugLogger.debug( + `Failed to get project root for worktree at ${cwd}: ${e instanceof Error ? e.message : String(e)}`, + ); + return cwd; + } +} + +export function getWorktreePath(projectRoot: string, name: string): string { + return path.join(projectRoot, '.gemini', 'worktrees', name); +} + +export async function createWorktree( + projectRoot: string, + name: string, +): Promise { + const worktreePath = getWorktreePath(projectRoot, name); + const branchName = `worktree-${name}`; + + await execa('git', ['worktree', 'add', worktreePath, '-b', branchName], { + cwd: projectRoot, + }); + + return worktreePath; +} + +export function isGeminiWorktree( + dirPath: string, + projectRoot: string, +): boolean { + try { + const realDirPath = realpathSync(dirPath); + const realProjectRoot = realpathSync(projectRoot); + const worktreesBaseDir = path.join(realProjectRoot, '.gemini', 'worktrees'); + const relative = path.relative(worktreesBaseDir, realDirPath); + return !relative.startsWith('..') && !path.isAbsolute(relative); + } catch { + return false; + } +} + +export async function hasWorktreeChanges( + dirPath: string, + baseSha?: string, +): Promise { + try { + // 1. Check for uncommitted changes (index or working tree) + const { stdout: status } = await execa('git', ['status', '--porcelain'], { + cwd: dirPath, + }); + if (status.trim() !== '') { + return true; + } + + // 2. Check if the current commit has moved from the base + if (baseSha) { + const { stdout: currentSha } = await execa('git', ['rev-parse', 'HEAD'], { + cwd: dirPath, + }); + if (currentSha.trim() !== baseSha) { + return true; + } + } + + return false; + } catch (e: unknown) { + debugLogger.debug( + `Failed to check worktree changes at ${dirPath}: ${e instanceof Error ? e.message : String(e)}`, + ); + // If any git command fails, assume the worktree is dirty to be safe. + return true; + } +} + +export async function cleanupWorktree( + dirPath: string, + projectRoot: string, +): Promise { + try { + await fs.access(dirPath); + } catch { + return; // Worktree already gone + } + + let branchName: string | undefined; + + try { + // 1. Discover the branch name associated with this worktree path + const { stdout } = await execa( + 'git', + ['-C', dirPath, 'branch', '--show-current'], + { + cwd: projectRoot, + }, + ); + branchName = stdout.trim() || undefined; + + // 2. Remove the worktree + await execa('git', ['worktree', 'remove', dirPath, '--force'], { + cwd: projectRoot, + }); + } catch (e: unknown) { + debugLogger.debug( + `Failed to remove worktree ${dirPath}: ${e instanceof Error ? e.message : String(e)}`, + ); + } finally { + // 3. Delete the branch if we found it + if (branchName) { + try { + await execa('git', ['branch', '-D', branchName], { + cwd: projectRoot, + }); + } catch (e: unknown) { + debugLogger.debug( + `Failed to delete branch ${branchName}: ${e instanceof Error ? e.message : String(e)}`, + ); + } + } + } +} diff --git a/packages/core/src/utils/memoryImportProcessor.ts b/packages/core/src/utils/memoryImportProcessor.ts index bf20bd6c13..10bf1ad592 100644 --- a/packages/core/src/utils/memoryImportProcessor.ts +++ b/packages/core/src/utils/memoryImportProcessor.ts @@ -48,16 +48,16 @@ export interface ProcessImportsResult { importTree: MemoryFile; } -// Helper to find the project root (looks for .git directory) +// Helper to find the project root (looks for .git directory or file for worktrees) async function findProjectRoot(startDir: string): Promise { let currentDir = path.resolve(startDir); while (true) { const gitPath = path.join(currentDir, '.git'); try { - const stats = await fs.lstat(gitPath); - if (stats.isDirectory()) { - return currentDir; - } + // Check for existence only — .git can be a directory (normal repos) + // or a file (submodules / worktrees). + await fs.access(gitPath); + return currentDir; } catch { // .git not found, continue to parent } diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 9c790c6268..85a907e57e 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -2663,6 +2663,13 @@ "default": true, "type": "boolean" }, + "worktrees": { + "title": "Enable Git Worktrees", + "description": "Enable automated Git worktree management for parallel work.", + "markdownDescription": "Enable automated Git worktree management for parallel work.\n\n- Category: `Experimental`\n- Requires restart: `yes`\n- Default: `false`", + "default": false, + "type": "boolean" + }, "extensionManagement": { "title": "Extension Management", "description": "Enable extension management features.", From 7a65c1e91dd483e6c0b4e6cb60cc19cdbc440840 Mon Sep 17 00:00:00 2001 From: kevinjwang1 Date: Fri, 20 Mar 2026 08:08:34 -0700 Subject: [PATCH 03/32] Add support for linking in the extension registry (#23153) --- .../src/ui/commands/extensionsCommand.test.ts | 24 ++++++--- .../cli/src/ui/commands/extensionsCommand.ts | 18 +++++-- .../views/ExtensionDetails.test.tsx | 51 +++++++++++++++++++ .../ui/components/views/ExtensionDetails.tsx | 27 +++++++++- .../views/ExtensionRegistryView.tsx | 24 +++++++++ 5 files changed, 131 insertions(+), 13 deletions(-) diff --git a/packages/cli/src/ui/commands/extensionsCommand.test.ts b/packages/cli/src/ui/commands/extensionsCommand.test.ts index dc49390c7e..8f065438e2 100644 --- a/packages/cli/src/ui/commands/extensionsCommand.test.ts +++ b/packages/cli/src/ui/commands/extensionsCommand.test.ts @@ -710,10 +710,14 @@ describe('extensionsCommand', () => { size: 100, } as Stats); await linkAction!(mockContext, packageName); - expect(mockInstallExtension).toHaveBeenCalledWith({ - source: packageName, - type: 'link', - }); + expect(mockInstallExtension).toHaveBeenCalledWith( + { + source: packageName, + type: 'link', + }, + undefined, + undefined, + ); expect(mockContext.ui.addItem).toHaveBeenCalledWith({ type: MessageType.INFO, text: `Linking extension from "${packageName}"...`, @@ -733,10 +737,14 @@ describe('extensionsCommand', () => { } as Stats); await linkAction!(mockContext, packageName); - expect(mockInstallExtension).toHaveBeenCalledWith({ - source: packageName, - type: 'link', - }); + expect(mockInstallExtension).toHaveBeenCalledWith( + { + source: packageName, + type: 'link', + }, + undefined, + undefined, + ); expect(mockContext.ui.addItem).toHaveBeenCalledWith({ type: MessageType.ERROR, text: `Failed to link extension from "${packageName}": ${errorMessage}`, diff --git a/packages/cli/src/ui/commands/extensionsCommand.ts b/packages/cli/src/ui/commands/extensionsCommand.ts index 8e988917e5..aed7595389 100644 --- a/packages/cli/src/ui/commands/extensionsCommand.ts +++ b/packages/cli/src/ui/commands/extensionsCommand.ts @@ -286,6 +286,11 @@ async function exploreAction( await installAction(context, extension.url, requestConsentOverride); context.ui.removeComponent(); }, + onLink: async (extension, requestConsentOverride) => { + debugLogger.log(`Linking extension: ${extension.extensionName}`); + await linkAction(context, extension.url, requestConsentOverride); + context.ui.removeComponent(); + }, onClose: () => context.ui.removeComponent(), extensionManager, }), @@ -533,7 +538,11 @@ async function installAction( } } -async function linkAction(context: CommandContext, args: string) { +async function linkAction( + context: CommandContext, + args: string, + requestConsentOverride?: (consent: string) => Promise, +) { const extensionLoader = context.services.agentContext?.config.getExtensionLoader(); if (!(extensionLoader instanceof ExtensionManager)) { @@ -582,8 +591,11 @@ async function linkAction(context: CommandContext, args: string) { source: sourceFilepath, type: 'link', }; - const extension = - await extensionLoader.installOrUpdateExtension(installMetadata); + const extension = await extensionLoader.installOrUpdateExtension( + installMetadata, + undefined, + requestConsentOverride, + ); context.ui.addItem({ type: MessageType.INFO, text: `Extension "${extension.name}" linked successfully.`, diff --git a/packages/cli/src/ui/components/views/ExtensionDetails.test.tsx b/packages/cli/src/ui/components/views/ExtensionDetails.test.tsx index 2da019d485..239f728472 100644 --- a/packages/cli/src/ui/components/views/ExtensionDetails.test.tsx +++ b/packages/cli/src/ui/components/views/ExtensionDetails.test.tsx @@ -32,13 +32,20 @@ const mockExtension: RegistryExtension = { licenseKey: 'Apache-2.0', }; +const linkableExtension: RegistryExtension = { + ...mockExtension, + url: '/local/path/to/extension', +}; + describe('ExtensionDetails', () => { let mockOnBack: ReturnType; let mockOnInstall: ReturnType; + let mockOnLink: ReturnType; beforeEach(() => { mockOnBack = vi.fn(); mockOnInstall = vi.fn(); + mockOnLink = vi.fn(); }); const renderDetails = async (isInstalled = false) => @@ -47,6 +54,7 @@ describe('ExtensionDetails', () => { extension={mockExtension} onBack={mockOnBack} onInstall={mockOnInstall} + onLink={mockOnLink} isInstalled={isInstalled} />, ); @@ -117,4 +125,47 @@ describe('ExtensionDetails', () => { expect(mockOnInstall).not.toHaveBeenCalled(); vi.useRealTimers(); }); + + it('should call onLink when "l" is pressed and is linkable', async () => { + const { stdin, waitUntilReady } = await renderWithProviders( + , + ); + await waitUntilReady(); + await React.act(async () => { + stdin.write('l'); + }); + await waitFor(() => { + expect(mockOnLink).toHaveBeenCalled(); + }); + }); + + it('should NOT show "Link" button for GitHub extensions', async () => { + const { lastFrame, waitUntilReady } = await renderDetails(false); + await waitUntilReady(); + await waitFor(() => { + expect(lastFrame()).not.toContain('[L] Link'); + }); + }); + + it('should show "Link" button for local extensions', async () => { + const { lastFrame, waitUntilReady } = await renderWithProviders( + , + ); + await waitUntilReady(); + await waitFor(() => { + expect(lastFrame()).toContain('[L] Link'); + }); + }); }); diff --git a/packages/cli/src/ui/components/views/ExtensionDetails.tsx b/packages/cli/src/ui/components/views/ExtensionDetails.tsx index 7ee38c0e54..82a6c42b78 100644 --- a/packages/cli/src/ui/components/views/ExtensionDetails.tsx +++ b/packages/cli/src/ui/components/views/ExtensionDetails.tsx @@ -19,6 +19,9 @@ export interface ExtensionDetailsProps { onInstall: ( requestConsentOverride: (consent: string) => Promise, ) => void | Promise; + onLink: ( + requestConsentOverride: (consent: string) => Promise, + ) => void | Promise; isInstalled: boolean; } @@ -26,6 +29,7 @@ export function ExtensionDetails({ extension, onBack, onInstall, + onLink, isInstalled, }: ExtensionDetailsProps): React.JSX.Element { const keyMatchers = useKeyMatchers(); @@ -35,6 +39,11 @@ export function ExtensionDetails({ } | null>(null); const [isInstalling, setIsInstalling] = useState(false); + const isLinkable = + !extension.url.startsWith('http') && + !extension.url.startsWith('git@') && + !extension.url.startsWith('sso://'); + useKeypress( (key) => { if (consentRequest) { @@ -56,6 +65,7 @@ export function ExtensionDetails({ onBack(); return true; } + if (keyMatchers[Command.RETURN](key) && !isInstalled && !isInstalling) { setIsInstalling(true); void onInstall( @@ -66,6 +76,16 @@ export function ExtensionDetails({ ); return true; } + if (key.name === 'l' && isLinkable && !isInstalled && !isInstalling) { + setIsInstalling(true); + void onLink( + (prompt: string) => + new Promise((resolve) => { + setConsentRequest({ prompt, resolve }); + }), + ); + return true; + } return false; }, { isActive: true, priority: true }, @@ -230,8 +250,11 @@ export function ExtensionDetails({ understand the permissions it requires and the actions it may perform. - - [{'Enter'}] Install + + + [{'Enter'}] Install + + {isLinkable && [L] Link} )} diff --git a/packages/cli/src/ui/components/views/ExtensionRegistryView.tsx b/packages/cli/src/ui/components/views/ExtensionRegistryView.tsx index 0539437fc3..60b0deec4a 100644 --- a/packages/cli/src/ui/components/views/ExtensionRegistryView.tsx +++ b/packages/cli/src/ui/components/views/ExtensionRegistryView.tsx @@ -29,6 +29,10 @@ export interface ExtensionRegistryViewProps { extension: RegistryExtension, requestConsentOverride?: (consent: string) => Promise, ) => void | Promise; + onLink?: ( + extension: RegistryExtension, + requestConsentOverride?: (consent: string) => Promise, + ) => void | Promise; onClose?: () => void; extensionManager: ExtensionManager; } @@ -39,6 +43,7 @@ interface ExtensionItem extends GenericListItem { export function ExtensionRegistryView({ onSelect, + onLink, onClose, extensionManager, }: ExtensionRegistryViewProps): React.JSX.Element { @@ -96,6 +101,22 @@ export function ExtensionRegistryView({ [onSelect, extensionManager], ); + const handleLink = useCallback( + async ( + extension: RegistryExtension, + requestConsentOverride?: (consent: string) => Promise, + ) => { + await onLink?.(extension, requestConsentOverride); + + // Refresh installed extensions list + setInstalledExtensions(extensionManager.getExtensions()); + + // Go back to the search page (list view) + setSelectedExtension(null); + }, + [onLink, extensionManager], + ); + const renderItem = useCallback( (item: ExtensionItem, isActive: boolean, _labelWidth: number) => { const isInstalled = installedExtensions.some( @@ -260,6 +281,9 @@ export function ExtensionRegistryView({ onInstall={async (requestConsentOverride) => { await handleInstall(selectedExtension, requestConsentOverride); }} + onLink={async (requestConsentOverride) => { + await handleLink(selectedExtension, requestConsentOverride); + }} isInstalled={installedExtensions.some( (e) => e.name === selectedExtension.extensionName, )} From 62cb14fa520f6f2caa29bd5bd7cf2b01afd15ab2 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Fri, 20 Mar 2026 21:40:59 +0530 Subject: [PATCH 04/32] feat(extensions): add --skip-settings flag to install command (#17212) --- docs/extensions/reference.md | 3 +- .../src/commands/extensions/install.test.ts | 124 +++++++++++------- .../cli/src/commands/extensions/install.ts | 10 +- 3 files changed, 86 insertions(+), 51 deletions(-) diff --git a/docs/extensions/reference.md b/docs/extensions/reference.md index 708caeb08d..56c51d30df 100644 --- a/docs/extensions/reference.md +++ b/docs/extensions/reference.md @@ -23,7 +23,7 @@ Gemini CLI creates a copy of the extension during installation. You must run GitHub, you must have `git` installed on your machine. ```bash -gemini extensions install [--ref ] [--auto-update] [--pre-release] [--consent] +gemini extensions install [--ref ] [--auto-update] [--pre-release] [--consent] [--skip-settings] ``` - ``: The GitHub URL or local path of the extension. @@ -31,6 +31,7 @@ gemini extensions install [--ref ] [--auto-update] [--pre-release] - `--auto-update`: Enable automatic updates for this extension. - `--pre-release`: Enable installation of pre-release versions. - `--consent`: Acknowledge security risks and skip the confirmation prompt. +- `--skip-settings`: Skip the configuration on install process. ### Uninstall an extension diff --git a/packages/cli/src/commands/extensions/install.test.ts b/packages/cli/src/commands/extensions/install.test.ts index 417e750651..8b3f8c5807 100644 --- a/packages/cli/src/commands/extensions/install.test.ts +++ b/packages/cli/src/commands/extensions/install.test.ts @@ -12,48 +12,46 @@ import { beforeEach, afterEach, type MockInstance, - type Mock, } from 'vitest'; import { handleInstall, installCommand } from './install.js'; import yargs from 'yargs'; import * as core from '@google/gemini-cli-core'; -import { - ExtensionManager, - type inferInstallMetadata, -} from '../../config/extension-manager.js'; -import type { - promptForConsentNonInteractive, - requestConsentNonInteractive, -} from '../../config/extensions/consent.js'; -import type { - isWorkspaceTrusted, - loadTrustedFolders, -} from '../../config/trustedFolders.js'; -import type * as fs from 'node:fs/promises'; import type { Stats } from 'node:fs'; import * as path from 'node:path'; +import { promptForSetting } from '../../config/extensions/extensionSettings.js'; -const mockInstallOrUpdateExtension: Mock< - typeof ExtensionManager.prototype.installOrUpdateExtension -> = vi.hoisted(() => vi.fn()); -const mockRequestConsentNonInteractive: Mock< - typeof requestConsentNonInteractive -> = vi.hoisted(() => vi.fn()); -const mockPromptForConsentNonInteractive: Mock< - typeof promptForConsentNonInteractive -> = vi.hoisted(() => vi.fn()); -const mockStat: Mock = vi.hoisted(() => vi.fn()); -const mockInferInstallMetadata: Mock = vi.hoisted( - () => vi.fn(), -); -const mockIsWorkspaceTrusted: Mock = vi.hoisted(() => - vi.fn(), -); -const mockLoadTrustedFolders: Mock = vi.hoisted(() => - vi.fn(), -); -const mockDiscover: Mock = - vi.hoisted(() => vi.fn()); +const { + mockInstallOrUpdateExtension, + mockLoadExtensions, + mockExtensionManager, + mockRequestConsentNonInteractive, + mockPromptForConsentNonInteractive, + mockStat, + mockInferInstallMetadata, + mockIsWorkspaceTrusted, + mockLoadTrustedFolders, + mockDiscover, +} = vi.hoisted(() => { + const mockLoadExtensions = vi.fn(); + const mockInstallOrUpdateExtension = vi.fn(); + const mockExtensionManager = vi.fn().mockImplementation(() => ({ + loadExtensions: mockLoadExtensions, + installOrUpdateExtension: mockInstallOrUpdateExtension, + })); + + return { + mockLoadExtensions, + mockInstallOrUpdateExtension, + mockExtensionManager, + mockRequestConsentNonInteractive: vi.fn(), + mockPromptForConsentNonInteractive: vi.fn(), + mockStat: vi.fn(), + mockInferInstallMetadata: vi.fn(), + mockIsWorkspaceTrusted: vi.fn(), + mockLoadTrustedFolders: vi.fn(), + mockDiscover: vi.fn(), + }; +}); vi.mock('../../config/extensions/consent.js', () => ({ requestConsentNonInteractive: mockRequestConsentNonInteractive, @@ -84,6 +82,7 @@ vi.mock('../../config/extension-manager.js', async (importOriginal) => ({ ...(await importOriginal< typeof import('../../config/extension-manager.js') >()), + ExtensionManager: mockExtensionManager, inferInstallMetadata: mockInferInstallMetadata, })); @@ -117,19 +116,18 @@ describe('handleInstall', () => { let processSpy: MockInstance; beforeEach(() => { - debugLogSpy = vi.spyOn(core.debugLogger, 'log'); - debugErrorSpy = vi.spyOn(core.debugLogger, 'error'); + debugLogSpy = vi + .spyOn(core.debugLogger, 'log') + .mockImplementation(() => {}); + debugErrorSpy = vi + .spyOn(core.debugLogger, 'error') + .mockImplementation(() => {}); processSpy = vi .spyOn(process, 'exit') .mockImplementation(() => undefined as never); - vi.spyOn(ExtensionManager.prototype, 'loadExtensions').mockResolvedValue( - [], - ); - vi.spyOn( - ExtensionManager.prototype, - 'installOrUpdateExtension', - ).mockImplementation(mockInstallOrUpdateExtension); + mockLoadExtensions.mockResolvedValue([]); + mockInstallOrUpdateExtension.mockReset(); mockIsWorkspaceTrusted.mockReturnValue({ isTrusted: true, source: 'file' }); mockDiscover.mockResolvedValue({ @@ -163,12 +161,7 @@ describe('handleInstall', () => { }); afterEach(() => { - mockInstallOrUpdateExtension.mockClear(); - mockRequestConsentNonInteractive.mockClear(); - mockStat.mockClear(); - mockInferInstallMetadata.mockClear(); vi.clearAllMocks(); - vi.restoreAllMocks(); }); function createMockExtension( @@ -288,6 +281,39 @@ describe('handleInstall', () => { expect(processSpy).toHaveBeenCalledWith(1); }); + it('should pass promptForSetting when skipSettings is not provided', async () => { + mockInstallOrUpdateExtension.mockResolvedValue({ + name: 'test-extension', + } as unknown as core.GeminiCLIExtension); + + await handleInstall({ + source: 'http://google.com', + }); + + expect(mockExtensionManager).toHaveBeenCalledWith( + expect.objectContaining({ + requestSetting: promptForSetting, + }), + ); + }); + + it('should pass null for requestSetting when skipSettings is true', async () => { + mockInstallOrUpdateExtension.mockResolvedValue({ + name: 'test-extension', + } as unknown as core.GeminiCLIExtension); + + await handleInstall({ + source: 'http://google.com', + skipSettings: true, + }); + + expect(mockExtensionManager).toHaveBeenCalledWith( + expect.objectContaining({ + requestSetting: null, + }), + ); + }); + it('should proceed if local path is already trusted', async () => { mockInstallOrUpdateExtension.mockResolvedValue( createMockExtension({ diff --git a/packages/cli/src/commands/extensions/install.ts b/packages/cli/src/commands/extensions/install.ts index 542d1240be..cf135a9366 100644 --- a/packages/cli/src/commands/extensions/install.ts +++ b/packages/cli/src/commands/extensions/install.ts @@ -37,6 +37,7 @@ interface InstallArgs { autoUpdate?: boolean; allowPreRelease?: boolean; consent?: boolean; + skipSettings?: boolean; } export async function handleInstall(args: InstallArgs) { @@ -153,7 +154,7 @@ export async function handleInstall(args: InstallArgs) { const extensionManager = new ExtensionManager({ workspaceDir, requestConsent, - requestSetting: promptForSetting, + requestSetting: args.skipSettings ? null : promptForSetting, settings, }); await extensionManager.loadExtensions(); @@ -196,6 +197,11 @@ export const installCommand: CommandModule = { type: 'boolean', default: false, }) + .option('skip-settings', { + describe: 'Skip the configuration on install process.', + type: 'boolean', + default: false, + }) .check((argv) => { if (!argv.source) { throw new Error('The source argument must be provided.'); @@ -214,6 +220,8 @@ export const installCommand: CommandModule = { allowPreRelease: argv['pre-release'] as boolean | undefined, // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion consent: argv['consent'] as boolean | undefined, + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + skipSettings: argv['skip-settings'] as boolean | undefined, }); await exitCli(); }, From b459e1a1082abf7779d5b236e10bb0c436b302b0 Mon Sep 17 00:00:00 2001 From: Jerop Kipruto Date: Fri, 20 Mar 2026 15:01:12 -0400 Subject: [PATCH 05/32] feat(telemetry): track if session is running in a Git worktree (#23265) --- docs/cli/telemetry.md | 1 + packages/cli/src/test-utils/mockConfig.ts | 1 + .../clearcut-logger/clearcut-logger.ts | 5 + .../clearcut-logger/event-metadata-key.ts | 3 + packages/core/src/telemetry/loggers.test.ts | 108 +++++++++++------- packages/core/src/telemetry/types.ts | 3 + 6 files changed, 80 insertions(+), 41 deletions(-) diff --git a/docs/cli/telemetry.md b/docs/cli/telemetry.md index 2068759213..fec0fb41c3 100644 --- a/docs/cli/telemetry.md +++ b/docs/cli/telemetry.md @@ -306,6 +306,7 @@ Emitted at startup with the CLI configuration. - `extension_ids` (string) - `extensions_count` (int) - `auth_type` (string) +- `worktree_active` (boolean) - `github_workflow_name` (string, optional) - `github_repository_hash` (string, optional) - `github_event_name` (string, optional) diff --git a/packages/cli/src/test-utils/mockConfig.ts b/packages/cli/src/test-utils/mockConfig.ts index d4f11212e3..e1505df970 100644 --- a/packages/cli/src/test-utils/mockConfig.ts +++ b/packages/cli/src/test-utils/mockConfig.ts @@ -44,6 +44,7 @@ export const createMockConfig = (overrides: Partial = {}): Config => getDeleteSession: vi.fn(() => undefined), setSessionId: vi.fn(), getSessionId: vi.fn().mockReturnValue('mock-session-id'), + getWorktreeSettings: vi.fn(() => undefined), getContentGeneratorConfig: vi.fn(() => ({ authType: 'google' })), getAcpMode: vi.fn(() => false), isBrowserLaunchSuppressed: vi.fn(() => false), diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts index 2f059030ca..11433db3e8 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts @@ -687,6 +687,11 @@ export class ClearcutLogger { gemini_cli_key: EventMetadataKey.GEMINI_CLI_START_SESSION_EXTENSION_IDS, value: event.extension_ids.toString(), }, + { + gemini_cli_key: + EventMetadataKey.GEMINI_CLI_START_SESSION_WORKTREE_ACTIVE, + value: event.worktree_active.toString(), + }, ]; // Add hardware information only to the start session event diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts index 632730aeeb..b7b9c0fd3a 100644 --- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts +++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts @@ -452,6 +452,9 @@ export enum EventMetadataKey { // Logs the name of extensions as a comma-separated string GEMINI_CLI_START_SESSION_EXTENSION_IDS = 120, + // Logs whether the session is running in a Git worktree. + GEMINI_CLI_START_SESSION_WORKTREE_ACTIVE = 191, + // Logs the setting scope for an extension enablement. GEMINI_CLI_EXTENSION_ENABLE_SETTING_SCOPE = 102, diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 4373a6b96c..27c23e7baa 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -195,48 +195,51 @@ describe('loggers', () => { }); describe('logCliConfiguration', () => { + const baseMockConfig = { + getSessionId: () => 'test-session-id', + getModel: () => 'test-model', + getEmbeddingModel: () => 'test-embedding-model', + getSandbox: () => true, + getCoreTools: () => ['ls', 'read-file'], + getApprovalMode: () => 'default', + getContentGeneratorConfig: () => ({ + model: 'test-model', + apiKey: 'test-api-key', + authType: AuthType.USE_VERTEX_AI, + }), + getTelemetryEnabled: () => true, + getUsageStatisticsEnabled: () => true, + getTelemetryLogPromptsEnabled: () => true, + getFileFilteringRespectGitIgnore: () => true, + getFileFilteringAllowBuildArtifacts: () => false, + getDebugMode: () => true, + getMcpServers: () => { + throw new Error('Should not call'); + }, + getQuestion: () => 'test-question', + getTargetDir: () => 'target-dir', + getProxy: () => 'http://test.proxy.com:8080', + getOutputFormat: () => OutputFormat.JSON, + getExtensions: () => + [ + { name: 'ext-one', id: 'id-one' }, + { name: 'ext-two', id: 'id-two' }, + ] as GeminiCLIExtension[], + getMcpClientManager: () => ({ + getMcpServers: () => ({ + 'test-server': { + command: 'test-command', + }, + }), + }), + isInteractive: () => false, + getExperiments: () => undefined, + getExperimentsAsync: async () => undefined, + getWorktreeSettings: () => undefined, + } as unknown as Config; + it('should log the cli configuration', async () => { - const mockConfig = { - getSessionId: () => 'test-session-id', - getModel: () => 'test-model', - getEmbeddingModel: () => 'test-embedding-model', - getSandbox: () => true, - getCoreTools: () => ['ls', 'read-file'], - getApprovalMode: () => 'default', - getContentGeneratorConfig: () => ({ - model: 'test-model', - apiKey: 'test-api-key', - authType: AuthType.USE_VERTEX_AI, - }), - getTelemetryEnabled: () => true, - getUsageStatisticsEnabled: () => true, - getTelemetryLogPromptsEnabled: () => true, - getFileFilteringRespectGitIgnore: () => true, - getFileFilteringAllowBuildArtifacts: () => false, - getDebugMode: () => true, - getMcpServers: () => { - throw new Error('Should not call'); - }, - getQuestion: () => 'test-question', - getTargetDir: () => 'target-dir', - getProxy: () => 'http://test.proxy.com:8080', - getOutputFormat: () => OutputFormat.JSON, - getExtensions: () => - [ - { name: 'ext-one', id: 'id-one' }, - { name: 'ext-two', id: 'id-two' }, - ] as GeminiCLIExtension[], - getMcpClientManager: () => ({ - getMcpServers: () => ({ - 'test-server': { - command: 'test-command', - }, - }), - }), - isInteractive: () => false, - getExperiments: () => undefined, - getExperimentsAsync: async () => undefined, - } as unknown as Config; + const mockConfig = baseMockConfig; const startSessionEvent = new StartSessionEvent(mockConfig); logCliConfiguration(mockConfig, startSessionEvent); @@ -270,9 +273,32 @@ describe('loggers', () => { extensions_count: 2, extensions: 'ext-one,ext-two', auth_type: 'vertex-ai', + worktree_active: false, }, }); }); + + it('should set worktree_active to true when worktree settings are present', async () => { + const mockConfig = { + ...baseMockConfig, + getWorktreeSettings: () => ({ + name: 'test-worktree', + path: '/path/to/worktree', + baseSha: 'test-sha', + }), + } as unknown as Config; + + const startSessionEvent = new StartSessionEvent(mockConfig); + logCliConfiguration(mockConfig, startSessionEvent); + + await new Promise(process.nextTick); + expect(mockLogger.emit).toHaveBeenCalledWith({ + body: 'CLI configuration loaded.', + attributes: expect.objectContaining({ + worktree_active: true, + }), + }); + }); }); describe('logUserPrompt', () => { diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index 0ee6e63503..1e0e3abc6e 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -77,6 +77,7 @@ export class StartSessionEvent implements BaseTelemetryEvent { extensions: string; extension_ids: string; auth_type?: string; + worktree_active: boolean; constructor(config: Config, toolRegistry?: ToolRegistry) { const generatorConfig = config.getContentGeneratorConfig(); @@ -114,6 +115,7 @@ export class StartSessionEvent implements BaseTelemetryEvent { this.extensions = extensions.map((e) => e.name).join(','); this.extension_ids = extensions.map((e) => e.id).join(','); this.auth_type = generatorConfig?.authType; + this.worktree_active = !!config.getWorktreeSettings(); if (toolRegistry) { const mcpTools = toolRegistry .getAllTools() @@ -147,6 +149,7 @@ export class StartSessionEvent implements BaseTelemetryEvent { extensions_count: this.extensions_count, extension_ids: this.extension_ids, auth_type: this.auth_type, + worktree_active: this.worktree_active, }; } From 26b9af1cdc433aa674e1f6cf6461939a02263c6e Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Fri, 20 Mar 2026 12:10:01 -0700 Subject: [PATCH 06/32] refactor(core): use absolute paths in GEMINI.md context markers (#23135) --- .../core/src/services/contextManager.test.ts | 2 +- packages/core/src/services/contextManager.ts | 8 +- .../core/src/utils/memoryDiscovery.test.ts | 74 +++++++++---------- packages/core/src/utils/memoryDiscovery.ts | 10 +-- 4 files changed, 40 insertions(+), 54 deletions(-) diff --git a/packages/core/src/services/contextManager.test.ts b/packages/core/src/services/contextManager.test.ts index 945c9263f6..1d078fd8fb 100644 --- a/packages/core/src/services/contextManager.test.ts +++ b/packages/core/src/services/contextManager.test.ts @@ -198,7 +198,7 @@ describe('ContextManager', () => { expect.any(Set), expect.any(Set), ); - expect(result).toMatch(/--- Context from: src[\\/]GEMINI\.md ---/); + expect(result).toMatch(/--- Context from: \/app\/src\/GEMINI\.md ---/); expect(result).toContain('Src Content'); expect(contextManager.getLoadedPaths()).toContain('/app/src/GEMINI.md'); }); diff --git a/packages/core/src/services/contextManager.ts b/packages/core/src/services/contextManager.ts index cec7c89ef9..b9da286e9c 100644 --- a/packages/core/src/services/contextManager.ts +++ b/packages/core/src/services/contextManager.ts @@ -98,12 +98,7 @@ export class ContextManager { paths: { global: string[]; extension: string[]; project: string[] }, contentsMap: Map, ) { - const workingDir = this.config.getWorkingDir(); - const hierarchicalMemory = categorizeAndConcatenate( - paths, - contentsMap, - workingDir, - ); + const hierarchicalMemory = categorizeAndConcatenate(paths, contentsMap); this.globalMemory = hierarchicalMemory.global || ''; this.extensionMemory = hierarchicalMemory.extension || ''; @@ -155,7 +150,6 @@ export class ContextManager { } return concatenateInstructions( result.files.map((f) => ({ filePath: f.path, content: f.content })), - this.config.getWorkingDir(), ); } diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts index f9c1671283..8ec6909b41 100644 --- a/packages/core/src/utils/memoryDiscovery.test.ts +++ b/packages/core/src/utils/memoryDiscovery.test.ts @@ -171,7 +171,7 @@ describe('memoryDiscovery', () => { ); expect(fileCount).toEqual(1); - expect(memoryContent).toContain(path.relative(cwd, filepath).toString()); + expect(memoryContent).toContain(filepath); expect(filePaths).toEqual([filepath]); }); }); @@ -215,9 +215,9 @@ describe('memoryDiscovery', () => { memoryContent: flattenMemory(result.memoryContent), }).toEqual({ memoryContent: `--- Global --- ---- Context from: ${path.relative(cwd, defaultContextFile)} --- +--- Context from: ${defaultContextFile} --- default context content ---- End of Context from: ${path.relative(cwd, defaultContextFile)} ---`, +--- End of Context from: ${defaultContextFile} ---`, fileCount: 1, filePaths: [defaultContextFile], }); @@ -244,9 +244,9 @@ default context content expect(result).toEqual({ memoryContent: `--- Global --- ---- Context from: ${normMarker(path.relative(cwd, customContextFile))} --- +--- Context from: ${customContextFile} --- custom context content ---- End of Context from: ${normMarker(path.relative(cwd, customContextFile))} ---`, +--- End of Context from: ${customContextFile} ---`, fileCount: 1, filePaths: [customContextFile], }); @@ -277,13 +277,13 @@ custom context content expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(path.relative(cwd, projectContextFile))} --- +--- Context from: ${projectContextFile} --- project context content ---- End of Context from: ${normMarker(path.relative(cwd, projectContextFile))} --- +--- End of Context from: ${projectContextFile} --- ---- Context from: ${normMarker(path.relative(cwd, cwdContextFile))} --- +--- Context from: ${cwdContextFile} --- cwd context content ---- End of Context from: ${normMarker(path.relative(cwd, cwdContextFile))} ---`, +--- End of Context from: ${cwdContextFile} ---`, fileCount: 2, filePaths: [projectContextFile, cwdContextFile], }); @@ -314,13 +314,13 @@ cwd context content expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(customFilename)} --- +--- Context from: ${cwdCustomFile} --- CWD custom memory ---- End of Context from: ${normMarker(customFilename)} --- +--- End of Context from: ${cwdCustomFile} --- ---- Context from: ${normMarker(path.join('subdir', customFilename))} --- +--- Context from: ${subdirCustomFile} --- Subdir custom memory ---- End of Context from: ${normMarker(path.join('subdir', customFilename))} ---`, +--- End of Context from: ${subdirCustomFile} ---`, fileCount: 2, filePaths: [cwdCustomFile, subdirCustomFile], }); @@ -348,13 +348,13 @@ Subdir custom memory expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(path.relative(cwd, projectRootGeminiFile))} --- +--- Context from: ${projectRootGeminiFile} --- Project root memory ---- End of Context from: ${normMarker(path.relative(cwd, projectRootGeminiFile))} --- +--- End of Context from: ${projectRootGeminiFile} --- ---- Context from: ${normMarker(path.relative(cwd, srcGeminiFile))} --- +--- Context from: ${srcGeminiFile} --- Src directory memory ---- End of Context from: ${normMarker(path.relative(cwd, srcGeminiFile))} ---`, +--- End of Context from: ${srcGeminiFile} ---`, fileCount: 2, filePaths: [projectRootGeminiFile, srcGeminiFile], }); @@ -382,13 +382,13 @@ Src directory memory expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(DEFAULT_CONTEXT_FILENAME)} --- +--- Context from: ${cwdGeminiFile} --- CWD memory ---- End of Context from: ${normMarker(DEFAULT_CONTEXT_FILENAME)} --- +--- End of Context from: ${cwdGeminiFile} --- ---- Context from: ${normMarker(path.join('subdir', DEFAULT_CONTEXT_FILENAME))} --- +--- Context from: ${subDirGeminiFile} --- Subdir memory ---- End of Context from: ${normMarker(path.join('subdir', DEFAULT_CONTEXT_FILENAME))} ---`, +--- End of Context from: ${subDirGeminiFile} ---`, fileCount: 2, filePaths: [cwdGeminiFile, subDirGeminiFile], }); @@ -428,26 +428,26 @@ Subdir memory expect(result).toEqual({ memoryContent: `--- Global --- ---- Context from: ${normMarker(path.relative(cwd, defaultContextFile))} --- +--- Context from: ${defaultContextFile} --- default context content ---- End of Context from: ${normMarker(path.relative(cwd, defaultContextFile))} --- +--- End of Context from: ${defaultContextFile} --- --- Project --- ---- Context from: ${normMarker(path.relative(cwd, rootGeminiFile))} --- +--- Context from: ${rootGeminiFile} --- Project parent memory ---- End of Context from: ${normMarker(path.relative(cwd, rootGeminiFile))} --- +--- End of Context from: ${rootGeminiFile} --- ---- Context from: ${normMarker(path.relative(cwd, projectRootGeminiFile))} --- +--- Context from: ${projectRootGeminiFile} --- Project root memory ---- End of Context from: ${normMarker(path.relative(cwd, projectRootGeminiFile))} --- +--- End of Context from: ${projectRootGeminiFile} --- ---- Context from: ${normMarker(path.relative(cwd, cwdGeminiFile))} --- +--- Context from: ${cwdGeminiFile} --- CWD memory ---- End of Context from: ${normMarker(path.relative(cwd, cwdGeminiFile))} --- +--- End of Context from: ${cwdGeminiFile} --- ---- Context from: ${normMarker(path.relative(cwd, subDirGeminiFile))} --- +--- Context from: ${subDirGeminiFile} --- Subdir memory ---- End of Context from: ${normMarker(path.relative(cwd, subDirGeminiFile))} ---`, +--- End of Context from: ${subDirGeminiFile} ---`, fileCount: 5, filePaths: [ defaultContextFile, @@ -491,9 +491,9 @@ Subdir memory expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(path.relative(cwd, regularSubDirGeminiFile))} --- +--- Context from: ${regularSubDirGeminiFile} --- My code memory ---- End of Context from: ${normMarker(path.relative(cwd, regularSubDirGeminiFile))} ---`, +--- End of Context from: ${regularSubDirGeminiFile} ---`, fileCount: 1, filePaths: [regularSubDirGeminiFile], }); @@ -565,9 +565,9 @@ My code memory expect(result).toEqual({ memoryContent: `--- Extension --- ---- Context from: ${normMarker(path.relative(cwd, extensionFilePath))} --- +--- Context from: ${extensionFilePath} --- Extension memory content ---- End of Context from: ${normMarker(path.relative(cwd, extensionFilePath))} ---`, +--- End of Context from: ${extensionFilePath} ---`, fileCount: 1, filePaths: [extensionFilePath], }); @@ -594,9 +594,9 @@ Extension memory content expect(result).toEqual({ memoryContent: `--- Project --- ---- Context from: ${normMarker(path.relative(cwd, includedFile))} --- +--- Context from: ${includedFile} --- included directory memory ---- End of Context from: ${normMarker(path.relative(cwd, includedFile))} ---`, +--- End of Context from: ${includedFile} ---`, fileCount: 1, filePaths: [includedFile], }); diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index 15b4b2c701..21b87330a1 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/packages/core/src/utils/memoryDiscovery.ts @@ -424,8 +424,6 @@ export async function readGeminiMdFiles( export function concatenateInstructions( instructionContents: GeminiFileContent[], - // CWD is needed to resolve relative paths for display markers - currentWorkingDirectoryForDisplay: string, ): string { return instructionContents .filter((item) => typeof item.content === 'string') @@ -435,10 +433,7 @@ export function concatenateInstructions( if (trimmedContent.length === 0) { return null; } - const displayPath = path.isAbsolute(item.filePath) - ? path.relative(currentWorkingDirectoryForDisplay, item.filePath) - : item.filePath; - return `--- Context from: ${displayPath} ---\n${trimmedContent}\n--- End of Context from: ${displayPath} ---`; + return `--- Context from: ${item.filePath} ---\n${trimmedContent}\n--- End of Context from: ${item.filePath} ---`; }) .filter((block): block is string => block !== null) .join('\n\n'); @@ -514,14 +509,12 @@ export async function getEnvironmentMemoryPaths( export function categorizeAndConcatenate( paths: { global: string[]; extension: string[]; project: string[] }, contentsMap: Map, - workingDir: string, ): HierarchicalMemory { const getConcatenated = (pList: string[]) => concatenateInstructions( pList .map((p) => contentsMap.get(p)) .filter((c): c is GeminiFileContent => !!c), - workingDir, ); return { @@ -687,7 +680,6 @@ export async function loadServerHierarchicalMemory( project: discoveryResult.project, }, contentsMap, - currentWorkingDirectory, ); return { From cca595971d51e38e58739674206621d9d4f22ddb Mon Sep 17 00:00:00 2001 From: Dev Randalpura Date: Fri, 20 Mar 2026 15:19:18 -0400 Subject: [PATCH 07/32] fix(core): add sanitization to sub agent thoughts and centralize utilities (#22828) --- .../browser/browserAgentInvocation.test.ts | 48 ++++++ .../agents/browser/browserAgentInvocation.ts | 137 +--------------- .../core/src/agents/local-invocation.test.ts | 33 ++++ packages/core/src/agents/local-invocation.ts | 23 ++- .../utils/agent-sanitization-utils.test.ts | 103 ++++++++++++ .../src/utils/agent-sanitization-utils.ts | 154 ++++++++++++++++++ 6 files changed, 362 insertions(+), 136 deletions(-) create mode 100644 packages/core/src/utils/agent-sanitization-utils.test.ts create mode 100644 packages/core/src/utils/agent-sanitization-utils.ts diff --git a/packages/core/src/agents/browser/browserAgentInvocation.test.ts b/packages/core/src/agents/browser/browserAgentInvocation.test.ts index 6cf47ae9d9..e41377bdd4 100644 --- a/packages/core/src/agents/browser/browserAgentInvocation.test.ts +++ b/packages/core/src/agents/browser/browserAgentInvocation.test.ts @@ -343,9 +343,57 @@ describe('BrowserAgentInvocation', () => { a.content.includes('Navigating to the page...'), ), ); + expect(thoughtProgress).toBeDefined(); }); + it('should overwrite the thought content with new THOUGHT_CHUNK activity', async () => { + const { fireActivity } = setupActivityCapture(); + const updateOutput = vi.fn(); + + const invocation = new BrowserAgentInvocation( + mockConfig, + mockParams, + mockMessageBus, + ); + + const executePromise = invocation.execute( + new AbortController().signal, + updateOutput, + ); + + // Allow createBrowserAgentDefinition to resolve and onActivity to be registered + await Promise.resolve(); + await Promise.resolve(); + + fireActivity({ + isSubagentActivityEvent: true, + agentName: 'browser_agent', + type: 'THOUGHT_CHUNK', + data: { text: 'I am thinking.' }, + }); + fireActivity({ + isSubagentActivityEvent: true, + agentName: 'browser_agent', + type: 'THOUGHT_CHUNK', + data: { text: 'Now I will act.' }, + }); + + await executePromise; + + const progressCalls = updateOutput.mock.calls + .map((c) => c[0] as SubagentProgress) + .filter((p) => p.isSubagentProgress); + + const lastCall = progressCalls[progressCalls.length - 1]; + expect(lastCall.recentActivity).toContainEqual( + expect.objectContaining({ + type: 'thought', + content: 'Now I will act.', + }), + ); + }); + it('should handle TOOL_CALL_START and TOOL_CALL_END with callId tracking', async () => { const { fireActivity } = setupActivityCapture(); const updateOutput = vi.fn(); diff --git a/packages/core/src/agents/browser/browserAgentInvocation.ts b/packages/core/src/agents/browser/browserAgentInvocation.ts index 5776aa85cd..60bd5201f0 100644 --- a/packages/core/src/agents/browser/browserAgentInvocation.ts +++ b/packages/core/src/agents/browser/browserAgentInvocation.ts @@ -37,138 +37,16 @@ import { cleanupBrowserAgent, } from './browserAgentFactory.js'; import { removeInputBlocker } from './inputBlocker.js'; +import { + sanitizeThoughtContent, + sanitizeToolArgs, + sanitizeErrorMessage, +} from '../../utils/agent-sanitization-utils.js'; const INPUT_PREVIEW_MAX_LENGTH = 50; const DESCRIPTION_MAX_LENGTH = 200; const MAX_RECENT_ACTIVITY = 20; -/** - * Sensitive key patterns used for redaction. - */ -const SENSITIVE_KEY_PATTERNS = [ - 'password', - 'pwd', - 'apikey', - 'api_key', - 'api-key', - 'token', - 'secret', - 'credential', - 'auth', - 'authorization', - 'access_token', - 'access_key', - 'refresh_token', - 'session_id', - 'cookie', - 'passphrase', - 'privatekey', - 'private_key', - 'private-key', - 'secret_key', - 'client_secret', - 'client_id', -]; - -/** - * Sanitizes tool arguments by recursively redacting sensitive fields. - * Supports nested objects and arrays. - */ -function sanitizeToolArgs(args: unknown): unknown { - if (typeof args === 'string') { - return sanitizeErrorMessage(args); - } - if (typeof args !== 'object' || args === null) { - return args; - } - - if (Array.isArray(args)) { - return args.map(sanitizeToolArgs); - } - - const sanitized: Record = {}; - - for (const [key, value] of Object.entries(args)) { - // Decode key to handle URL-encoded sensitive keys (e.g., api%5fkey) - let decodedKey = key; - try { - decodedKey = decodeURIComponent(key); - } catch { - // Ignore decoding errors - } - const keyNormalized = decodedKey.toLowerCase().replace(/[-_]/g, ''); - const isSensitive = SENSITIVE_KEY_PATTERNS.some((pattern) => - keyNormalized.includes(pattern.replace(/[-_]/g, '')), - ); - if (isSensitive) { - sanitized[key] = '[REDACTED]'; - } else { - sanitized[key] = sanitizeToolArgs(value); - } - } - - return sanitized; -} - -/** - * Sanitizes error messages by redacting potential sensitive data patterns. - * Uses [^\s'"]+ to catch JWTs, tokens with dots/slashes, and other complex values. - */ -function sanitizeErrorMessage(message: string): string { - if (!message) return message; - - let sanitized = message; - - // 1. Redact inline PEM content - sanitized = sanitized.replace( - /-----BEGIN\s+[\w\s]+-----[\s\S]*?-----END\s+[\w\s]+-----/g, - '[REDACTED_PEM]', - ); - - const unquotedValue = `[^\\s]+(?:\\s+(?![a-zA-Z0-9_.-]+(?:=|:))[^\\s=:<>]+)*`; - const valuePattern = `(?:"[^"]*"|'[^']*'|${unquotedValue})`; - - // 2. Handle key-value pairs with delimiters (=, :, space, CLI-style --flag) - const urlSafeKeyPatternStr = SENSITIVE_KEY_PATTERNS.map((p) => - p.replace(/[-_]/g, '(?:[-_]|%2D|%5F|%2d|%5f)?'), - ).join('|'); - - const keyWithDelimiter = new RegExp( - `((?:--)?("|')?(${urlSafeKeyPatternStr})\\2\\s*(?:[:=]|%3A|%3D)\\s*)${valuePattern}`, - 'gi', - ); - sanitized = sanitized.replace(keyWithDelimiter, '$1[REDACTED]'); - - // 3. Handle space-separated sensitive keywords (e.g. "password mypass", "--api-key secret") - const tokenValuePattern = `[A-Za-z0-9._\\-/+=]{8,}`; - const spaceKeywords = [ - ...SENSITIVE_KEY_PATTERNS.map((p) => - p.replace(/[-_]/g, '(?:[-_]|%2D|%5F|%2d|%5f)?'), - ), - 'bearer', - ]; - const spaceSeparated = new RegExp( - `\\b((?:--)?(?:${spaceKeywords.join('|')})(?:\\s*:\\s*bearer)?\\s+)(${tokenValuePattern})`, - 'gi', - ); - sanitized = sanitized.replace(spaceSeparated, '$1[REDACTED]'); - - // 4. Handle file path redaction - sanitized = sanitized.replace( - /((?:[/\\][a-zA-Z0-9_-]+)*[/\\][a-zA-Z0-9_-]*\.(?:key|pem|p12|pfx))/gi, - '/path/to/[REDACTED].key', - ); - - return sanitized; -} - -/** - * Sanitizes LLM thought content by redacting sensitive data patterns. - */ -function sanitizeThoughtContent(text: string): string { - return sanitizeErrorMessage(text); -} - /** * Browser agent invocation with async tool setup. * @@ -284,14 +162,13 @@ export class BrowserAgentInvocation extends BaseToolInvocation< case 'THOUGHT_CHUNK': { const text = String(activity.data['text']); const lastItem = recentActivity[recentActivity.length - 1]; + if ( lastItem && lastItem.type === 'thought' && lastItem.status === 'running' ) { - lastItem.content = sanitizeThoughtContent( - lastItem.content + text, - ); + lastItem.content = sanitizeThoughtContent(text); } else { recentActivity.push({ id: randomUUID(), diff --git a/packages/core/src/agents/local-invocation.test.ts b/packages/core/src/agents/local-invocation.test.ts index 34df9844c9..2153f538c9 100644 --- a/packages/core/src/agents/local-invocation.test.ts +++ b/packages/core/src/agents/local-invocation.test.ts @@ -271,6 +271,39 @@ describe('LocalSubagentInvocation', () => { ); }); + it('should overwrite the thought content with new THOUGHT_CHUNK activity', async () => { + mockExecutorInstance.run.mockImplementation(async () => { + const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2]; + + if (onActivity) { + onActivity({ + isSubagentActivityEvent: true, + agentName: 'MockAgent', + type: 'THOUGHT_CHUNK', + data: { text: 'I am thinking.' }, + } as SubagentActivityEvent); + onActivity({ + isSubagentActivityEvent: true, + agentName: 'MockAgent', + type: 'THOUGHT_CHUNK', + data: { text: 'Now I will act.' }, + } as SubagentActivityEvent); + } + return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL }; + }); + + await invocation.execute(signal, updateOutput); + + const calls = updateOutput.mock.calls; + const lastCall = calls[calls.length - 1][0] as SubagentProgress; + expect(lastCall.recentActivity).toContainEqual( + expect.objectContaining({ + type: 'thought', + content: 'Now I will act.', + }), + ); + }); + it('should stream other activities (e.g., TOOL_CALL_START, ERROR)', async () => { mockExecutorInstance.run.mockImplementation(async () => { const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2]; diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index e8b98d4744..08a4aa8264 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -24,6 +24,11 @@ import { } from './types.js'; import { randomUUID } from 'node:crypto'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { + sanitizeThoughtContent, + sanitizeToolArgs, + sanitizeErrorMessage, +} from '../utils/agent-sanitization-utils.js'; const INPUT_PREVIEW_MAX_LENGTH = 50; const DESCRIPTION_MAX_LENGTH = 200; @@ -118,17 +123,18 @@ export class LocalSubagentInvocation extends BaseToolInvocation< case 'THOUGHT_CHUNK': { const text = String(activity.data['text']); const lastItem = recentActivity[recentActivity.length - 1]; + if ( lastItem && lastItem.type === 'thought' && lastItem.status === 'running' ) { - lastItem.content = text; + lastItem.content = sanitizeThoughtContent(text); } else { recentActivity.push({ id: randomUUID(), type: 'thought', - content: text, + content: sanitizeThoughtContent(text), status: 'running', }); } @@ -138,12 +144,14 @@ export class LocalSubagentInvocation extends BaseToolInvocation< case 'TOOL_CALL_START': { const name = String(activity.data['name']); const displayName = activity.data['displayName'] - ? String(activity.data['displayName']) + ? sanitizeErrorMessage(String(activity.data['displayName'])) : undefined; const description = activity.data['description'] - ? String(activity.data['description']) + ? sanitizeErrorMessage(String(activity.data['description'])) : undefined; - const args = JSON.stringify(activity.data['args']); + const args = JSON.stringify( + sanitizeToolArgs(activity.data['args']), + ); recentActivity.push({ id: randomUUID(), type: 'tool_call', @@ -175,6 +183,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< case 'ERROR': { const error = String(activity.data['error']); const errorType = activity.data['errorType']; + const sanitizedError = sanitizeErrorMessage(error); const isCancellation = errorType === SubagentActivityErrorType.CANCELLED || error === SUBAGENT_CANCELLED_ERROR_MESSAGE; @@ -217,7 +226,9 @@ export class LocalSubagentInvocation extends BaseToolInvocation< id: randomUUID(), type: 'thought', content: - isCancellation || isRejection ? error : `Error: ${error}`, + isCancellation || isRejection + ? sanitizedError + : `Error: ${sanitizedError}`, status: isCancellation || isRejection ? 'cancelled' : 'error', }); updated = true; diff --git a/packages/core/src/utils/agent-sanitization-utils.test.ts b/packages/core/src/utils/agent-sanitization-utils.test.ts new file mode 100644 index 0000000000..fa030024a6 --- /dev/null +++ b/packages/core/src/utils/agent-sanitization-utils.test.ts @@ -0,0 +1,103 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { + sanitizeErrorMessage, + sanitizeToolArgs, + sanitizeThoughtContent, +} from './agent-sanitization-utils.js'; + +describe('agent-sanitization-utils', () => { + describe('sanitizeErrorMessage', () => { + it('should redact standard inline PEM content', () => { + const input = + 'Here is my key: -----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA12345\n-----END RSA PRIVATE KEY----- do not share.'; + const expected = 'Here is my key: [REDACTED_PEM] do not share.'; + expect(sanitizeErrorMessage(input)).toBe(expected); + }); + + it('should redact non-standard inline PEM content (with punctuation)', () => { + const input = + '-----BEGIN X.509 CERTIFICATE-----\nMIIEowIBAAKCAQEA12345\n-----END X.509 CERTIFICATE-----'; + const expected = '[REDACTED_PEM]'; + expect(sanitizeErrorMessage(input)).toBe(expected); + }); + + it('should not hang on ReDoS attack string for PEM redaction', () => { + const start = Date.now(); + // A string that starts with -----BEGIN but has no ending, with many spaces + // In the vulnerable regex, this would cause catastrophic backtracking. + const maliciousInput = '-----BEGIN ' + ' '.repeat(50000) + 'A'; + const result = sanitizeErrorMessage(maliciousInput); + const duration = Date.now() - start; + + // Should process very quickly (e.g. < 50ms) + expect(duration).toBeLessThan(50); + + // Since it doesn't match the full PEM block pattern, it should return the input unaltered + expect(result).toBe(maliciousInput); + }); + + it('should redact key-value pairs with sensitive keys', () => { + const input = 'Error: connection failed. --api-key="secret123"'; + const result = sanitizeErrorMessage(input); + expect(result).toContain('[REDACTED]'); + expect(result).not.toContain('secret123'); + }); + + it('should redact space-separated sensitive keywords', () => { + // The keyword regex requires tokens to be 8+ chars + const input = 'Using password mySuperSecretPassword123'; + const result = sanitizeErrorMessage(input); + expect(result).toContain('[REDACTED]'); + expect(result).not.toContain('mySuperSecretPassword123'); + }); + }); + + describe('sanitizeToolArgs', () => { + it('should redact sensitive fields in an object', () => { + const input = { + username: 'admin', + password: 'superSecretPassword', + nested: { + api_key: 'abc123xyz', + normal_field: 'hello', + }, + }; + + const result = sanitizeToolArgs(input); + + expect(result).toEqual({ + username: 'admin', + password: '[REDACTED]', + nested: { + api_key: '[REDACTED]', + normal_field: 'hello', + }, + }); + }); + + it('should handle arrays and strings correctly', () => { + const input = ['normal string', '--api-key="secret123"']; + const result = sanitizeToolArgs(input) as string[]; + + expect(result[0]).toBe('normal string'); + expect(result[1]).toContain('[REDACTED]'); + expect(result[1]).not.toContain('secret123'); + }); + }); + + describe('sanitizeThoughtContent', () => { + it('should redact sensitive patterns from thought content', () => { + const input = 'I will now authenticate using token 1234567890abcdef.'; + const result = sanitizeThoughtContent(input); + + expect(result).toContain('[REDACTED]'); + expect(result).not.toContain('1234567890abcdef'); + }); + }); +}); diff --git a/packages/core/src/utils/agent-sanitization-utils.ts b/packages/core/src/utils/agent-sanitization-utils.ts new file mode 100644 index 0000000000..e83c879fae --- /dev/null +++ b/packages/core/src/utils/agent-sanitization-utils.ts @@ -0,0 +1,154 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Sensitive key patterns used for redaction. + */ +export const SENSITIVE_KEY_PATTERNS = [ + 'password', + 'pwd', + 'apikey', + 'api_key', + 'api-key', + 'token', + 'secret', + 'credential', + 'auth', + 'authorization', + 'access_token', + 'access_key', + 'refresh_token', + 'session_id', + 'cookie', + 'passphrase', + 'privatekey', + 'private_key', + 'private-key', + 'secret_key', + 'client_secret', + 'client_id', +]; + +/** + * Sanitizes tool arguments by recursively redacting sensitive fields. + * Supports nested objects and arrays. + */ +export function sanitizeToolArgs(args: unknown): unknown { + if (typeof args === 'string') { + return sanitizeErrorMessage(args); + } + if (typeof args !== 'object' || args === null) { + return args; + } + + if (Array.isArray(args)) { + return args.map(sanitizeToolArgs); + } + + const sanitized: Record = {}; + + for (const [key, value] of Object.entries(args)) { + // Decode key to handle URL-encoded sensitive keys (e.g., api%5fkey) + let decodedKey = key; + try { + decodedKey = decodeURIComponent(key); + } catch { + // Ignore decoding errors + } + const keyNormalized = decodedKey.toLowerCase().replace(/[-_]/g, ''); + const isSensitive = SENSITIVE_KEY_PATTERNS.some((pattern) => + keyNormalized.includes(pattern.replace(/[-_]/g, '')), + ); + if (isSensitive) { + sanitized[key] = '[REDACTED]'; + } else { + sanitized[key] = sanitizeToolArgs(value); + } + } + + return sanitized; +} + +/** + * Sanitizes error messages by redacting potential sensitive data patterns. + * Uses [^\s'"]+ to catch JWTs, tokens with dots/slashes, and other complex values. + */ +export function sanitizeErrorMessage(message: string): string { + if (!message) return message; + + let sanitized = message; + + // 1. Redact inline PEM content (Safe iterative approach to avoid ReDoS) + let startIndex = 0; + while ((startIndex = sanitized.indexOf('-----BEGIN', startIndex)) !== -1) { + const endOfBegin = sanitized.indexOf('-----', startIndex + 10); + if (endOfBegin === -1) { + break; // No closing dashes for the BEGIN header + } + + // Find the END header + const endHeaderStart = sanitized.indexOf('-----END', endOfBegin + 5); + if (endHeaderStart === -1) { + break; // No END header found + } + + const endHeaderEnd = sanitized.indexOf('-----', endHeaderStart + 8); + if (endHeaderEnd === -1) { + break; // No closing dashes for the END header + } + + // We found a complete block. Replace it. + const before = sanitized.substring(0, startIndex); + const after = sanitized.substring(endHeaderEnd + 5); + sanitized = before + '[REDACTED_PEM]' + after; + + // Resume searching after the redacted block + startIndex = before.length + 14; // length of '[REDACTED_PEM]' + } + + const unquotedValue = `[^\\s]+(?:\\s+(?![a-zA-Z0-9_.-]+(?:=|:))[^\\s=:<>]+)*`; + const valuePattern = `(?:"[^"]*"|'[^']*'|${unquotedValue})`; + + // 2. Handle key-value pairs with delimiters (=, :, space, CLI-style --flag) + const urlSafeKeyPatternStr = SENSITIVE_KEY_PATTERNS.map((p) => + p.replace(/[-_]/g, '(?:[-_]|%2D|%5F|%2d|%5f)?'), + ).join('|'); + + const keyWithDelimiter = new RegExp( + `((?:--)?("|')?(${urlSafeKeyPatternStr})\\2\\s*(?:[:=]|%3A|%3D)\\s*)${valuePattern}`, + 'gi', + ); + sanitized = sanitized.replace(keyWithDelimiter, '$1[REDACTED]'); + + // 3. Handle space-separated sensitive keywords (e.g. "password mypass", "--api-key secret") + const tokenValuePattern = `[A-Za-z0-9._\\-/+=]{8,}`; + const spaceKeywords = [ + ...SENSITIVE_KEY_PATTERNS.map((p) => + p.replace(/[-_]/g, '(?:[-_]|%2D|%5F|%2d|%5f)?'), + ), + 'bearer', + ]; + const spaceSeparated = new RegExp( + `\\b((?:--)?(?:${spaceKeywords.join('|')})(?:\\s*:\\s*bearer)?\\s+)(${tokenValuePattern})`, + 'gi', + ); + sanitized = sanitized.replace(spaceSeparated, '$1[REDACTED]'); + + // 4. Handle file path redaction + sanitized = sanitized.replace( + /((?:[/\\][a-zA-Z0-9_-]+)*[/\\][a-zA-Z0-9_-]*\.(?:key|pem|p12|pfx))/gi, + '/path/to/[REDACTED].key', + ); + + return sanitized; +} + +/** + * Sanitizes LLM thought content by redacting sensitive data patterns. + */ +export function sanitizeThoughtContent(text: string): string { + return sanitizeErrorMessage(text); +} From 05e4ea80eed76be095af0bbdbdd63c16fc738f6b Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Fri, 20 Mar 2026 15:31:01 -0400 Subject: [PATCH 08/32] feat(core): refine User-Agent for VS Code traffic (unified format) (#23256) --- .../core/src/core/contentGenerator.test.ts | 121 +++++++++++++++++- packages/core/src/core/contentGenerator.ts | 39 +++++- packages/core/src/utils/surface.ts | 7 +- 3 files changed, 156 insertions(+), 11 deletions(-) diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 4bacd1b488..a264b2fb6c 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -131,6 +131,10 @@ describe('createContentGenerator', () => { // Set a fixed version for testing vi.stubEnv('CLI_VERSION', '1.2.3'); + vi.stubEnv('TERM_PROGRAM', 'iTerm.app'); + vi.stubEnv('VSCODE_PID', ''); + vi.stubEnv('GITHUB_SHA', ''); + vi.stubEnv('GEMINI_CLI_SURFACE', ''); const mockGenerator = { models: {}, @@ -149,7 +153,7 @@ describe('createContentGenerator', () => { httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.stringMatching( - /GeminiCLI\/1\.2\.3\/gemini-pro \(.*; .*; .*\)/, + /GeminiCLI\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/, ), }), }), @@ -159,7 +163,7 @@ describe('createContentGenerator', () => { ); }); - it('should include clientName prefix in User-Agent when specified', async () => { + it('should use standard User-Agent for a2a-server running outside VS Code', async () => { const mockConfig = { getModel: vi.fn().mockReturnValue('gemini-pro'), getProxy: vi.fn().mockReturnValue(undefined), @@ -169,6 +173,10 @@ describe('createContentGenerator', () => { // Set a fixed version for testing vi.stubEnv('CLI_VERSION', '1.2.3'); + vi.stubEnv('TERM_PROGRAM', 'iTerm.app'); + vi.stubEnv('VSCODE_PID', ''); + vi.stubEnv('GITHUB_SHA', ''); + vi.stubEnv('GEMINI_CLI_SURFACE', ''); const mockGenerator = { models: {}, @@ -185,7 +193,7 @@ describe('createContentGenerator', () => { httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.stringMatching( - /GeminiCLI-a2a-server\/.*\/gemini-pro \(.*; .*; .*\)/, + /GeminiCLI-a2a-server\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/, ), }), }), @@ -193,6 +201,113 @@ describe('createContentGenerator', () => { ); }); + it('should include unified User-Agent for a2a-server (VS Code Agent Mode)', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => true, + getClientName: vi.fn().mockReturnValue('a2a-server'), + } as unknown as Config; + + // Set a fixed version for testing + vi.stubEnv('CLI_VERSION', '1.2.3'); + // Mock the environment variable that the VS Code extension host would provide to the a2a-server process + vi.stubEnv('VSCODE_PID', '12345'); + vi.stubEnv('TERM_PROGRAM', 'vscode'); + vi.stubEnv('TERM_PROGRAM_VERSION', '1.85.0'); + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + await createContentGenerator( + { apiKey: 'test-api-key', authType: AuthType.USE_GEMINI }, + mockConfig, + undefined, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + httpOptions: expect.objectContaining({ + headers: expect.objectContaining({ + 'User-Agent': expect.stringMatching( + /CloudCodeVSCode\/1\.2\.3 \(aidev_client; os_type=.*; os_version=.*; arch=.*; host_path=VSCode\/1\.85\.0; proxy_client=geminicli\)/, + ), + }), + }), + }), + ); + }); + + it('should include clientName prefix in User-Agent when specified (non-VSCode)', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => true, + getClientName: vi.fn().mockReturnValue('my-client'), + } as unknown as Config; + + // Set a fixed version for testing + vi.stubEnv('CLI_VERSION', '1.2.3'); + vi.stubEnv('TERM_PROGRAM', 'iTerm.app'); + vi.stubEnv('VSCODE_PID', ''); + vi.stubEnv('GITHUB_SHA', ''); + vi.stubEnv('GEMINI_CLI_SURFACE', ''); + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + await createContentGenerator( + { apiKey: 'test-api-key', authType: AuthType.USE_GEMINI }, + mockConfig, + undefined, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + httpOptions: expect.objectContaining({ + headers: expect.objectContaining({ + 'User-Agent': expect.stringMatching( + /GeminiCLI-my-client\/1\.2\.3\/gemini-pro \(.*; .*; terminal\)/, + ), + }), + }), + }), + ); + }); + + it('should allow custom headers to override User-Agent', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => true, + getClientName: vi.fn().mockReturnValue(undefined), + } as unknown as Config; + + vi.stubEnv('GEMINI_CLI_CUSTOM_HEADERS', 'User-Agent:MyCustomUA'); + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + await createContentGenerator( + { apiKey: 'test-api-key', authType: AuthType.USE_GEMINI }, + mockConfig, + undefined, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + httpOptions: expect.objectContaining({ + headers: expect.objectContaining({ + 'User-Agent': 'MyCustomUA', + }), + }), + }), + ); + }); + it('should include custom headers from GEMINI_CLI_CUSTOM_HEADERS for Code Assist requests', async () => { const mockGenerator = {} as unknown as ContentGenerator; vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index ff1739c04b..c901562eb7 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -13,7 +13,9 @@ import { type EmbedContentResponse, type EmbedContentParameters, } from '@google/genai'; +import * as os from 'node:os'; import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js'; +import { isCloudShell } from '../ide/detect-ide.js'; import type { Config } from '../config/config.js'; import { loadApiKey } from './apiKeyCredentialStorage.js'; @@ -185,19 +187,46 @@ export async function createContentGenerator( const customHeadersEnv = process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined; const clientName = gcConfig.getClientName(); - const userAgentPrefix = clientName - ? `GeminiCLI-${clientName}` - : 'GeminiCLI'; const surface = determineSurface(); - const userAgent = `${userAgentPrefix}/${version}/${model} (${process.platform}; ${process.arch}; ${surface})`; + + let userAgent: string; + // Use unified format for VS Code traffic. + // Note: We don't automatically assume a2a-server is VS Code, + // as it could be used by other clients unless the surface explicitly says 'vscode'. + if (clientName === 'acp-vscode' || surface === 'vscode') { + const osTypeMap: Record = { + darwin: 'macOS', + win32: 'Windows', + linux: 'Linux', + }; + const osType = osTypeMap[process.platform] || process.platform; + const osVersion = os.release(); + const arch = process.arch; + + const vscodeVersion = process.env['TERM_PROGRAM_VERSION'] || 'unknown'; + let hostPath = `VSCode/${vscodeVersion}`; + if (isCloudShell()) { + const cloudShellVersion = + process.env['CLOUD_SHELL_VERSION'] || 'unknown'; + hostPath += ` > CloudShell/${cloudShellVersion}`; + } + + userAgent = `CloudCodeVSCode/${version} (aidev_client; os_type=${osType}; os_version=${osVersion}; arch=${arch}; host_path=${hostPath}; proxy_client=geminicli)`; + } else { + const userAgentPrefix = clientName + ? `GeminiCLI-${clientName}` + : 'GeminiCLI'; + userAgent = `${userAgentPrefix}/${version}/${model} (${process.platform}; ${process.arch}; ${surface})`; + } + const customHeadersMap = parseCustomHeaders(customHeadersEnv); const apiKeyAuthMechanism = process.env['GEMINI_API_KEY_AUTH_MECHANISM'] || 'x-goog-api-key'; const apiVersionEnv = process.env['GOOGLE_GENAI_API_VERSION']; const baseHeaders: Record = { - ...customHeadersMap, 'User-Agent': userAgent, + ...customHeadersMap, }; if ( diff --git a/packages/core/src/utils/surface.ts b/packages/core/src/utils/surface.ts index e4b1241d84..7c6bd4da6b 100644 --- a/packages/core/src/utils/surface.ts +++ b/packages/core/src/utils/surface.ts @@ -37,9 +37,10 @@ export function determineSurface(): string { return ide.name; } - // If the detected IDE is 'vscode', we only accept it if TERM_PROGRAM confirms it. - // This prevents generic terminals from being misidentified as VSCode. - if (process.env['TERM_PROGRAM'] === 'vscode') { + // If the detected IDE is 'vscode', we only accept it if TERM_PROGRAM or VSCODE_PID confirms it. + // This prevents generic terminals from being misidentified as VSCode, while still detecting + // background processes spawned by the VS Code extension host (like a2a-server). + if (process.env['TERM_PROGRAM'] === 'vscode' || process.env['VSCODE_PID']) { return ide.name; } From 86a3a913b5840dde19b5079e9a00dd4aa4142c0c Mon Sep 17 00:00:00 2001 From: kevinjwang1 Date: Fri, 20 Mar 2026 12:50:15 -0700 Subject: [PATCH 09/32] Fix schema for ModelChains (#23284) --- .../cli/src/config/settingsSchema.test.ts | 24 +++++++++++++++++++ packages/cli/src/config/settingsSchema.ts | 10 +++++++- .../cli/src/ui/components/ModelDialog.tsx | 3 ++- schemas/settings.schema.json | 10 +++++++- 4 files changed, 44 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/config/settingsSchema.test.ts b/packages/cli/src/config/settingsSchema.test.ts index 37ddf87642..c358cd65aa 100644 --- a/packages/cli/src/config/settingsSchema.test.ts +++ b/packages/cli/src/config/settingsSchema.test.ts @@ -538,8 +538,32 @@ describe('SettingsSchema', () => { } }; + const visitJsonSchema = (jsonSchema: Record) => { + const ref = jsonSchema['ref']; + if (typeof ref === 'string') { + referenced.add(ref); + } + const properties = jsonSchema['properties']; + if ( + properties && + typeof properties === 'object' && + !Array.isArray(properties) + ) { + Object.values(properties as Record).forEach((prop) => + visitJsonSchema(prop as Record), + ); + } + const items = jsonSchema['items']; + if (items && typeof items === 'object' && !Array.isArray(items)) { + visitJsonSchema(items as Record); + } + }; + Object.values(schema).forEach(visitDefinition); + // Also visit all definitions to find nested references + Object.values(SETTINGS_SCHEMA_DEFINITIONS).forEach(visitJsonSchema); + // Ensure definitions map doesn't accumulate stale entries. Object.keys(SETTINGS_SCHEMA_DEFINITIONS).forEach((key) => { if (!referenced.has(key)) { diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 3724253e97..3a622460aa 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1094,7 +1094,7 @@ const SETTINGS_SCHEMA = { showInDialog: false, additionalProperties: { type: 'array', - ref: 'ModelPolicy', + ref: 'ModelPolicyChain', }, }, }, @@ -2998,6 +2998,14 @@ export const SETTINGS_SCHEMA_DEFINITIONS: Record< }, }, }, + ModelPolicyChain: { + type: 'array', + description: 'A chain of model policies for fallback behavior.', + items: { + type: 'object', + ref: 'ModelPolicy', + }, + }, ModelPolicy: { type: 'object', description: diff --git a/packages/cli/src/ui/components/ModelDialog.tsx b/packages/cli/src/ui/components/ModelDialog.tsx index 85cf16de3b..c42838c070 100644 --- a/packages/cli/src/ui/components/ModelDialog.tsx +++ b/packages/cli/src/ui/components/ModelDialog.tsx @@ -233,7 +233,8 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element { }); // Deduplicate: only show one entry per unique resolved model value. - // This is needed because 3 pro and 3.1 pro models can resolve to the same value. + // This is needed because 3 pro and 3.1 pro models can resolve to the same + // value, depending on the useGemini31 flag. const seen = new Set(); return list.filter((option) => { if (seen.has(option.value)) return false; diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 85a907e57e..a231558bf7 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -2061,7 +2061,7 @@ }, "type": "object", "additionalProperties": { - "$ref": "#/$defs/ModelPolicy" + "$ref": "#/$defs/ModelPolicyChain" } } }, @@ -3686,6 +3686,14 @@ } } }, + "ModelPolicyChain": { + "type": "array", + "description": "A chain of model policies for fallback behavior.", + "items": { + "type": "object", + "ref": "ModelPolicy" + } + }, "ModelPolicy": { "type": "object", "description": "Defines the policy for a single model in the availability chain.", From 6c78eb7a39cc2fee281d50e245ebac8f259ed0a7 Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Fri, 20 Mar 2026 20:08:29 +0000 Subject: [PATCH 10/32] test(cli): refactor tests for async render utilities (#23252) --- .../cli/src/config/extensions/consent.test.ts | 5 +- packages/cli/src/test-utils/render.test.tsx | 53 +- packages/cli/src/test-utils/render.tsx | 63 +- packages/cli/src/ui/App.test.tsx | 134 +-- packages/cli/src/ui/AppContainer.test.tsx | 942 +++++++----------- .../cli/src/ui/IdeIntegrationNudge.test.tsx | 13 +- .../cli/src/ui/auth/ApiAuthDialog.test.tsx | 15 +- packages/cli/src/ui/auth/AuthDialog.test.tsx | 75 +- .../cli/src/ui/auth/AuthInProgress.test.tsx | 17 +- .../src/ui/auth/BannedAccountDialog.test.tsx | 30 +- .../LoginWithGoogleRestartDialog.test.tsx | 9 +- packages/cli/src/ui/auth/useAuth.test.tsx | 203 ++-- .../cli/src/ui/components/AboutBox.test.tsx | 12 +- .../AdminSettingsChangedDialog.test.tsx | 9 +- .../ui/components/AgentConfigDialog.test.tsx | 1 - .../AlternateBufferQuittingDisplay.test.tsx | 18 +- .../cli/src/ui/components/AnsiOutput.test.tsx | 24 +- .../cli/src/ui/components/AppHeader.test.tsx | 26 +- .../components/ApprovalModeIndicator.test.tsx | 18 +- .../src/ui/components/AskUserDialog.test.tsx | 25 +- .../BackgroundShellDisplay.test.tsx | 36 +- .../cli/src/ui/components/Checklist.test.tsx | 15 +- .../src/ui/components/ChecklistItem.test.tsx | 9 +- .../cli/src/ui/components/CliSpinner.test.tsx | 13 +- .../src/ui/components/ColorsDisplay.test.tsx | 3 +- .../cli/src/ui/components/Composer.test.tsx | 3 +- .../ui/components/ConfigInitDisplay.test.tsx | 5 +- .../src/ui/components/ConsentPrompt.test.tsx | 15 +- .../components/ConsoleSummaryDisplay.test.tsx | 6 +- .../components/ContextSummaryDisplay.test.tsx | 3 +- .../components/ContextUsageDisplay.test.tsx | 15 +- .../ui/components/CopyModeWarning.test.tsx | 6 +- .../src/ui/components/DebugProfiler.test.tsx | 12 +- .../DetailedMessagesDisplay.test.tsx | 15 +- .../src/ui/components/DialogManager.test.tsx | 6 +- .../components/EditorSettingsDialog.test.tsx | 11 +- .../ui/components/EmptyWalletDialog.test.tsx | 33 +- .../ui/components/ExitPlanModeDialog.test.tsx | 60 +- .../src/ui/components/ExitWarning.test.tsx | 12 +- .../ui/components/FolderTrustDialog.test.tsx | 51 +- .../cli/src/ui/components/Footer.test.tsx | 728 ++++++-------- .../ui/components/FooterConfigDialog.test.tsx | 23 +- .../GeminiRespondingSpinner.test.tsx | 21 +- .../ui/components/GradientRegression.test.tsx | 28 +- .../cli/src/ui/components/Header.test.tsx | 22 +- packages/cli/src/ui/components/Help.test.tsx | 9 +- .../ui/components/HistoryItemDisplay.test.tsx | 157 ++- .../ui/components/HookStatusDisplay.test.tsx | 12 +- .../src/ui/components/HooksDialog.test.tsx | 28 +- .../components/IdeTrustChangeDialog.test.tsx | 12 +- .../ui/components/LoadingIndicator.test.tsx | 52 +- .../LogoutConfirmationDialog.test.tsx | 9 +- .../LoopDetectionConfirmation.test.tsx | 6 +- .../src/ui/components/MainContent.test.tsx | 73 +- .../ui/components/MemoryUsageDisplay.test.tsx | 8 +- .../src/ui/components/ModelDialog.test.tsx | 1 - .../ui/components/ModelStatsDisplay.test.tsx | 6 +- .../MultiFolderTrustDialog.test.tsx | 24 +- .../components/NewAgentsNotification.test.tsx | 6 +- .../src/ui/components/Notifications.test.tsx | 39 +- .../ui/components/OverageMenuDialog.test.tsx | 44 +- .../PermissionsModifyTrustDialog.test.tsx | 15 +- .../ui/components/PolicyUpdateDialog.test.tsx | 3 +- .../src/ui/components/ProQuotaDialog.test.tsx | 40 +- .../components/QueuedMessageDisplay.test.tsx | 15 +- .../ui/components/QuittingDisplay.test.tsx | 6 +- .../src/ui/components/QuotaDisplay.test.tsx | 27 +- .../components/RawMarkdownIndicator.test.tsx | 10 +- .../ui/components/RewindConfirmation.test.tsx | 12 +- .../src/ui/components/RewindViewer.test.tsx | 38 +- .../src/ui/components/SessionBrowser.test.tsx | 18 +- .../SessionBrowserSearchNav.test.tsx | 23 +- .../SessionBrowserStates.test.tsx | 9 +- .../src/ui/components/SettingsDialog.test.tsx | 146 +-- .../ui/components/ShellInputPrompt.test.tsx | 27 +- .../ui/components/ShellModeIndicator.test.tsx | 5 +- .../src/ui/components/ShortcutsHelp.test.tsx | 5 +- .../src/ui/components/ShowMoreLines.test.tsx | 15 +- .../components/ShowMoreLinesLayout.test.tsx | 6 +- .../src/ui/components/StatsDisplay.test.tsx | 60 +- .../src/ui/components/StatusDisplay.test.tsx | 3 +- .../src/ui/components/StickyHeader.test.tsx | 3 +- .../ui/components/SuggestionsDisplay.test.tsx | 21 +- packages/cli/src/ui/components/Table.test.tsx | 27 +- .../src/ui/components/ThemeDialog.test.tsx | 21 +- .../src/ui/components/ThemedGradient.test.tsx | 3 +- packages/cli/src/ui/components/Tips.test.tsx | 5 +- .../src/ui/components/ToastDisplay.test.tsx | 30 +- .../components/ToolConfirmationQueue.test.tsx | 46 +- .../ui/components/ToolStatsDisplay.test.tsx | 3 +- .../ui/components/UpdateNotification.test.tsx | 3 +- .../src/ui/components/UserIdentity.test.tsx | 21 +- .../ui/components/ValidationDialog.test.tsx | 24 +- .../ConfigInitDisplay.test.tsx.snap | 12 + .../messages/CompressionMessage.test.tsx | 42 +- .../components/messages/ErrorMessage.test.tsx | 6 +- .../messages/GeminiMessage.test.tsx | 9 +- .../components/messages/InfoMessage.test.tsx | 11 +- .../messages/RedirectionConfirmation.test.tsx | 3 +- .../messages/ShellToolMessage.test.tsx | 53 +- .../messages/SubagentGroupDisplay.test.tsx | 6 +- .../messages/SubagentProgressDisplay.test.tsx | 24 +- .../src/ui/components/messages/Todo.test.tsx | 6 +- .../messages/ToolConfirmationMessage.test.tsx | 96 +- .../messages/ToolGroupMessage.test.tsx | 102 +- .../components/messages/ToolMessage.test.tsx | 72 +- .../messages/ToolMessageFocusHint.test.tsx | 3 - .../messages/ToolMessageRawMarkdown.test.tsx | 3 +- .../components/messages/ToolShared.test.tsx | 15 +- .../components/messages/UserMessage.test.tsx | 12 +- .../messages/WarningMessage.test.tsx | 6 +- .../shared/BaseSelectionList.test.tsx | 2 - .../components/shared/EnumSelector.test.tsx | 25 +- .../components/shared/ExpandableText.test.tsx | 35 +- .../shared/HalfLinePaddedBox.test.tsx | 12 +- .../ui/components/shared/MaxSizedBox.test.tsx | 24 +- .../ui/components/shared/Scrollable.test.tsx | 45 +- .../components/shared/SearchableList.test.tsx | 9 +- .../components/shared/SectionHeader.test.tsx | 3 +- .../shared/SlicingMaxSizedBox.test.tsx | 15 +- .../ui/components/shared/TabHeader.test.tsx | 39 +- .../ui/components/shared/TextInput.test.tsx | 39 +- .../shared/VirtualizedList.test.tsx | 31 +- .../ui/components/shared/performance.test.ts | 8 +- .../ui/components/shared/text-buffer.test.ts | 522 +++++----- .../src/ui/components/views/ChatList.test.tsx | 13 +- .../views/ExtensionDetails.test.tsx | 9 +- .../views/ExtensionRegistryView.test.tsx | 3 +- .../components/views/ExtensionsList.test.tsx | 15 +- .../ui/components/views/McpStatus.test.tsx | 52 +- .../ui/components/views/SkillsList.test.tsx | 18 +- .../ui/components/views/ToolsList.test.tsx | 9 +- .../ui/contexts/ScrollProvider.drag.test.tsx | 12 +- .../src/ui/contexts/ScrollProvider.test.tsx | 22 +- .../src/ui/contexts/SessionContext.test.tsx | 20 +- .../src/ui/contexts/SettingsContext.test.tsx | 20 +- .../src/ui/contexts/TerminalContext.test.tsx | 6 +- .../ui/contexts/ToolActionsContext.test.tsx | 40 +- .../ui/hooks/shellCommandProcessor.test.tsx | 74 +- .../ui/hooks/slashCommandProcessor.test.tsx | 2 +- .../src/ui/hooks/useAlternateBuffer.test.ts | 12 +- .../ui/hooks/useAnimatedScrollbar.test.tsx | 26 +- .../ui/hooks/useApprovalModeIndicator.test.ts | 84 +- .../cli/src/ui/hooks/useAtCompletion.test.ts | 122 ++- .../ui/hooks/useAtCompletion_agents.test.ts | 4 +- .../hooks/useBackgroundShellManager.test.tsx | 28 +- packages/cli/src/ui/hooks/useBanner.test.ts | 20 +- .../cli/src/ui/hooks/useBatchedScroll.test.ts | 28 +- .../src/ui/hooks/useConsoleMessages.test.tsx | 24 +- .../src/ui/hooks/useEditorSettings.test.tsx | 40 +- .../src/ui/hooks/useExtensionUpdates.test.tsx | 8 +- .../src/ui/hooks/useFlickerDetector.test.ts | 26 +- .../cli/src/ui/hooks/useFolderTrust.test.ts | 40 +- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 4 + .../src/ui/hooks/useGitBranchName.test.tsx | 143 +-- .../src/ui/hooks/useHistoryManager.test.ts | 44 +- .../src/ui/hooks/useHookDisplayState.test.ts | 24 +- .../src/ui/hooks/useIdeTrustListener.test.tsx | 29 +- .../src/ui/hooks/useIncludeDirsTrust.test.tsx | 18 +- .../src/ui/hooks/useInlineEditBuffer.test.ts | 36 +- .../cli/src/ui/hooks/useInputHistory.test.ts | 56 +- .../src/ui/hooks/useInputHistoryStore.test.ts | 40 +- .../src/ui/hooks/useLoadingIndicator.test.tsx | 36 +- packages/cli/src/ui/hooks/useLogger.test.tsx | 33 +- .../cli/src/ui/hooks/useMcpStatus.test.tsx | 20 +- .../src/ui/hooks/useMemoryMonitor.test.tsx | 12 +- .../cli/src/ui/hooks/useMessageQueue.test.tsx | 58 +- .../cli/src/ui/hooks/useModelCommand.test.tsx | 12 +- packages/cli/src/ui/hooks/useMouse.test.ts | 16 +- .../cli/src/ui/hooks/useMouseClick.test.ts | 4 +- .../hooks/usePermissionsModifyTrust.test.ts | 34 +- .../cli/src/ui/hooks/usePhraseCycler.test.tsx | 41 +- .../src/ui/hooks/usePrivacySettings.test.tsx | 34 +- .../src/ui/hooks/useQuotaAndFallback.test.ts | 62 +- packages/cli/src/ui/hooks/useRewind.test.ts | 20 +- .../src/ui/hooks/useSelectionList.test.tsx | 6 +- .../src/ui/hooks/useSessionBrowser.test.ts | 6 +- .../cli/src/ui/hooks/useSessionResume.test.ts | 46 +- .../ui/hooks/useSettingsNavigation.test.ts | 32 +- .../cli/src/ui/hooks/useShellHistory.test.ts | 14 +- .../ui/hooks/useShellInactivityStatus.test.ts | 14 +- .../src/ui/hooks/useSlashCompletion.test.ts | 533 +++++----- packages/cli/src/ui/hooks/useSuspend.test.ts | 12 +- .../src/ui/hooks/useTabbedNavigation.test.ts | 104 +- .../src/ui/hooks/useTerminalTheme.test.tsx | 33 +- packages/cli/src/ui/hooks/useTimer.test.tsx | 36 +- .../cli/src/ui/hooks/useToolScheduler.test.ts | 42 +- .../ui/hooks/useTurnActivityMonitor.test.ts | 19 +- .../cli/src/ui/hooks/vim-passthrough.test.tsx | 4 +- packages/cli/src/ui/hooks/vim.test.tsx | 498 ++++----- .../src/ui/layouts/DefaultAppLayout.test.tsx | 9 +- .../privacy/CloudFreePrivacyNotice.test.tsx | 9 +- .../privacy/CloudPaidPrivacyNotice.test.tsx | 6 +- .../ui/privacy/GeminiPrivacyNotice.test.tsx | 6 +- .../cli/src/ui/privacy/PrivacyNotice.test.tsx | 3 +- .../cli/src/ui/utils/CodeColorizer.test.tsx | 5 +- .../cli/src/ui/utils/MarkdownDisplay.test.tsx | 48 +- .../cli/src/ui/utils/TableRenderer.test.tsx | 47 +- 198 files changed, 3592 insertions(+), 4802 deletions(-) diff --git a/packages/cli/src/config/extensions/consent.test.ts b/packages/cli/src/config/extensions/consent.test.ts index 76d7227ab4..8de884cdd5 100644 --- a/packages/cli/src/config/extensions/consent.test.ts +++ b/packages/cli/src/config/extensions/consent.test.ts @@ -59,8 +59,9 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }); async function expectConsentSnapshot(consentString: string) { - const renderResult = render(React.createElement(Text, null, consentString)); - await renderResult.waitUntilReady(); + const renderResult = await render( + React.createElement(Text, null, consentString), + ); await expect(renderResult).toMatchSvgSnapshot(); } diff --git a/packages/cli/src/test-utils/render.test.tsx b/packages/cli/src/test-utils/render.test.tsx index 7172a99119..3c3f4102a4 100644 --- a/packages/cli/src/test-utils/render.test.tsx +++ b/packages/cli/src/test-utils/render.test.tsx @@ -12,24 +12,18 @@ import { waitFor } from './async.js'; describe('render', () => { it('should render a component', async () => { - const { lastFrame, waitUntilReady, unmount } = render( - Hello World, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await render(Hello World); expect(lastFrame()).toBe('Hello World\n'); unmount(); }); it('should support rerender', async () => { - const { lastFrame, rerender, waitUntilReady, unmount } = render( + const { lastFrame, rerender, waitUntilReady, unmount } = await render( Hello, ); - await waitUntilReady(); expect(lastFrame()).toBe('Hello\n'); - await act(async () => { - rerender(World); - }); + await act(async () => rerender(World)); await waitUntilReady(); expect(lastFrame()).toBe('World\n'); unmount(); @@ -42,10 +36,8 @@ describe('render', () => { return Hello; } - const { unmount, waitUntilReady } = render(); - await waitUntilReady(); + const { unmount } = await render(); unmount(); - expect(cleanupMock).toHaveBeenCalled(); }); }); @@ -54,36 +46,27 @@ describe('renderHook', () => { it('should rerender with previous props when called without arguments', async () => { const useTestHook = ({ value }: { value: number }) => { const [count, setCount] = useState(0); - useEffect(() => { - setCount((c) => c + 1); - }, [value]); + useEffect(() => setCount((c) => c + 1), [value]); return { count, value }; }; - const { result, rerender, waitUntilReady, unmount } = renderHook( + const { result, rerender, waitUntilReady, unmount } = await renderHook( useTestHook, - { - initialProps: { value: 1 }, - }, + { initialProps: { value: 1 } }, ); - await waitUntilReady(); expect(result.current.value).toBe(1); await waitFor(() => expect(result.current.count).toBe(1)); // Rerender with new props - await act(async () => { - rerender({ value: 2 }); - }); + await act(async () => rerender({ value: 2 })); await waitUntilReady(); expect(result.current.value).toBe(2); await waitFor(() => expect(result.current.count).toBe(2)); // Rerender without arguments should use previous props (value: 2) // This would previously crash or pass undefined if not fixed - await act(async () => { - rerender(); - }); + await act(async () => rerender()); await waitUntilReady(); expect(result.current.value).toBe(2); // Count should not increase because value didn't change @@ -98,14 +81,11 @@ describe('renderHook', () => { }; const { result, rerender, waitUntilReady, unmount } = - renderHook(useTestHook); - await waitUntilReady(); + await renderHook(useTestHook); expect(result.current.count).toBe(0); - await act(async () => { - rerender(); - }); + await act(async () => rerender()); await waitUntilReady(); expect(result.current.count).toBe(0); unmount(); @@ -113,19 +93,14 @@ describe('renderHook', () => { it('should update props if undefined is passed explicitly', async () => { const useTestHook = (val: string | undefined) => val; - const { result, rerender, waitUntilReady, unmount } = renderHook( + const { result, rerender, waitUntilReady, unmount } = await renderHook( useTestHook, - { - initialProps: 'initial' as string | undefined, - }, + { initialProps: 'initial' }, ); - await waitUntilReady(); expect(result.current).toBe('initial'); - await act(async () => { - rerender(undefined); - }); + await act(async () => rerender(undefined)); await waitUntilReady(); expect(result.current).toBeUndefined(); unmount(); diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index 7d298b120d..ea889181c6 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -257,13 +257,9 @@ class XtermStdout extends EventEmitter { return currentFrame !== ''; } - // If both are empty, it's a match. - // We consider undefined lastRenderOutput as effectively empty for this check - // to support hook testing where Ink may skip rendering completely. - if ( - (this.lastRenderOutput === undefined || expectedFrame === '') && - currentFrame === '' - ) { + // If Ink expects nothing (no new static content and no dynamic output), + // we consider it a match because the terminal buffer will just hold the historical static content. + if (expectedFrame === '') { return true; } @@ -271,8 +267,8 @@ class XtermStdout extends EventEmitter { return false; } - // If Ink expects nothing but terminal has content, or vice-versa, it's NOT a match. - if (expectedFrame === '' || currentFrame === '') { + // If the terminal is empty but Ink expects something, it's not a match. + if (currentFrame === '') { return false; } @@ -382,13 +378,11 @@ export type RenderInstance = { const instances: InkInstance[] = []; -// Wrapper around ink's render that ensures act() is called and uses Xterm for output -export const render = ( +export const render = async ( tree: React.ReactElement, terminalWidth?: number, -): Omit< - RenderInstance, - 'capturedOverflowState' | 'capturedOverflowActions' +): Promise< + Omit > => { const cols = terminalWidth ?? 100; // We use 1000 rows to avoid windows with incorrect snapshots if a correct @@ -437,6 +431,8 @@ export const render = ( instances.push(instance); + await stdout.waitUntilReady(); + return { rerender: (newTree: React.ReactElement) => { act(() => { @@ -751,7 +747,10 @@ export const renderWithProviders = async ( ); - const renderResult = render(wrapWithProviders(component), terminalWidth); + const renderResult = await render( + wrapWithProviders(component), + terminalWidth, + ); return { ...renderResult, @@ -765,19 +764,19 @@ export const renderWithProviders = async ( }; }; -export function renderHook( +export async function renderHook( renderCallback: (props: Props) => Result, options?: { initialProps?: Props; wrapper?: React.ComponentType<{ children: React.ReactNode }>; }, -): { +): Promise<{ result: { current: Result }; rerender: (props?: Props) => void; unmount: () => void; waitUntilReady: () => Promise; generateSvg: () => string; -} { +}> { const result = { current: undefined as unknown as Result }; let currentProps = options?.initialProps as Props; @@ -800,17 +799,15 @@ export function renderHook( let waitUntilReady: () => Promise = async () => {}; let generateSvg: () => string = () => ''; - act(() => { - const renderResult = render( - - - , - ); - inkRerender = renderResult.rerender; - unmount = renderResult.unmount; - waitUntilReady = renderResult.waitUntilReady; - generateSvg = renderResult.generateSvg; - }); + const renderResult = await render( + + + , + ); + inkRerender = renderResult.rerender; + unmount = renderResult.unmount; + waitUntilReady = renderResult.waitUntilReady; + generateSvg = renderResult.generateSvg; function rerender(props?: Props) { if (arguments.length > 0) { @@ -864,7 +861,13 @@ export async function renderHookWithProviders( const Wrapper = options.wrapper || (({ children }) => <>{children}); - let renderResult: ReturnType; + let renderResult: RenderInstance & { + simulateClick: ( + col: number, + row: number, + button?: 0 | 1 | 2, + ) => Promise; + }; await act(async () => { renderResult = await renderWithProviders( diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 7f5e55c022..950363f6a8 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -94,14 +94,10 @@ describe('App', () => { }; it('should render main content and composer when not quitting', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: mockUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: false } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: mockUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: false } }), + }); expect(lastFrame()).toContain('Tips for getting started'); expect(lastFrame()).toContain('Notifications'); @@ -115,14 +111,10 @@ describe('App', () => { quittingMessages: [{ id: 1, type: 'user', text: 'test' }], } as UIState; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: quittingUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: false } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: quittingUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: false } }), + }); expect(lastFrame()).toContain('Quitting...'); unmount(); @@ -136,14 +128,10 @@ describe('App', () => { pendingHistoryItems: [{ type: 'user', text: 'pending item' }], } as UIState; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: quittingUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: quittingUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain('HistoryItemDisplay'); expect(lastFrame()).toContain('Quitting...'); @@ -156,14 +144,10 @@ describe('App', () => { dialogsVisible: true, } as UIState; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: dialogUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: dialogUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain('Tips for getting started'); expect(lastFrame()).toContain('Notifications'); @@ -183,14 +167,10 @@ describe('App', () => { [stateKey]: true, } as UIState; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain(`Press Ctrl+${key} again to exit.`); unmount(); @@ -200,14 +180,10 @@ describe('App', () => { it('should render ScreenReaderAppLayout when screen reader is enabled', async () => { (useIsScreenReaderEnabled as Mock).mockReturnValue(true); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: mockUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: mockUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain('Notifications'); expect(lastFrame()).toContain('Footer'); @@ -219,14 +195,10 @@ describe('App', () => { it('should render DefaultAppLayout when screen reader is not enabled', async () => { (useIsScreenReaderEnabled as Mock).mockReturnValue(false); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: mockUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: mockUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain('Tips for getting started'); expect(lastFrame()).toContain('Notifications'); @@ -274,15 +246,11 @@ describe('App', () => { vi.spyOn(configWithExperiment, 'isTrustedFolder').mockReturnValue(true); vi.spyOn(configWithExperiment, 'getIdeMode').mockReturnValue(false); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: stateWithConfirmingTool, - config: configWithExperiment, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: stateWithConfirmingTool, + config: configWithExperiment, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toContain('Tips for getting started'); expect(lastFrame()).toContain('Notifications'); @@ -295,28 +263,20 @@ describe('App', () => { describe('Snapshots', () => { it('renders default layout correctly', async () => { (useIsScreenReaderEnabled as Mock).mockReturnValue(false); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: mockUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: mockUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('renders screen reader layout correctly', async () => { (useIsScreenReaderEnabled as Mock).mockReturnValue(true); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: mockUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: mockUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toMatchSnapshot(); unmount(); }); @@ -326,14 +286,10 @@ describe('App', () => { ...mockUIState, dialogsVisible: true, } as UIState; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { - uiState: dialogUIState, - settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + uiState: dialogUIState, + settings: createMockSettings({ ui: { useAlternateBuffer: true } }), + }); expect(lastFrame()).toMatchSnapshot(); unmount(); }); diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 650804025b..313573a573 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -16,7 +16,7 @@ import { } from 'vitest'; import { render, cleanup, persistentStateMock } from '../test-utils/render.js'; import { waitFor } from '../test-utils/async.js'; -import { act, useContext, type ReactElement } from 'react'; +import { act, useContext } from 'react'; import { AppContainer } from './AppContainer.js'; import { SettingsContext } from './contexts/SettingsContext.js'; import { type TrackedToolCall } from './hooks/useToolScheduler.js'; @@ -250,6 +250,15 @@ describe('AppContainer State Management', () => { let mockInitResult: InitializationResult; let mockExtensionManager: MockedObject; + type AppContainerProps = { + settings?: LoadedSettings; + config?: Config; + version?: string; + initResult?: InitializationResult; + startupWarnings?: StartupWarning[]; + resumedSessionData?: ResumedSessionData; + }; + // Helper to generate the AppContainer JSX for render and rerender const getAppContainer = ({ settings = mockSettings, @@ -258,14 +267,7 @@ describe('AppContainer State Management', () => { initResult = mockInitResult, startupWarnings, resumedSessionData, - }: { - settings?: LoadedSettings; - config?: Config; - version?: string; - initResult?: InitializationResult; - startupWarnings?: StartupWarning[]; - resumedSessionData?: ResumedSessionData; - } = {}) => ( + }: AppContainerProps = {}) => ( @@ -282,7 +284,7 @@ describe('AppContainer State Management', () => { ); // Helper to render the AppContainer - const renderAppContainer = (props?: Parameters[0]) => + const renderAppContainer = async (props?: AppContainerProps) => render(getAppContainer(props)); // Create typed mocks for all hooks @@ -514,13 +516,9 @@ describe('AppContainer State Management', () => { describe('Basic Rendering', () => { it('renders without crashing with minimal props', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + expect(capturedUIState).toBeTruthy(); + unmount(); }); it('renders with startup warnings', async () => { @@ -537,44 +535,32 @@ describe('AppContainer State Management', () => { }, ]; - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ startupWarnings }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => + renderAppContainer({ startupWarnings }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }); it('shows full UI details by default', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => { - expect(capturedUIState.cleanUiDetailsVisible).toBe(true); - }); - unmount!(); + expect(capturedUIState.cleanUiDetailsVisible).toBe(true); + unmount(); }); it('starts in minimal UI mode when Focus UI preference is persisted', async () => { persistentStateMock.get.mockReturnValueOnce(true); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ settings: mockSettings, - }); - unmount = result.unmount; - }); + }), + ); - await waitFor(() => { - expect(capturedUIState.cleanUiDetailsVisible).toBe(false); - }); + expect(capturedUIState.cleanUiDetailsVisible).toBe(false); expect(persistentStateMock.get).toHaveBeenCalledWith('focusUiEnabled'); - unmount!(); + unmount(); }); }); @@ -609,15 +595,9 @@ describe('AppContainer State Management', () => { ], }); - let unmount: (() => void) | undefined; - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => - expect(terminalNotificationsMocks.notifyViaTerminal).toHaveBeenCalled(), - ); + expect(terminalNotificationsMocks.notifyViaTerminal).toHaveBeenCalled(); expect( terminalNotificationsMocks.buildRunEventNotificationContent, ).toHaveBeenCalledWith( @@ -626,9 +606,7 @@ describe('AppContainer State Management', () => { }), ); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('does not send attention notification when terminal is focused', async () => { @@ -661,19 +639,13 @@ describe('AppContainer State Management', () => { ], }); - let unmount: (() => void) | undefined; - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); expect( terminalNotificationsMocks.notifyViaTerminal, ).not.toHaveBeenCalled(); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('sends attention notification when focus reporting is unavailable', async () => { @@ -706,19 +678,11 @@ describe('AppContainer State Management', () => { ], }); - let unmount: (() => void) | undefined; - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => - expect(terminalNotificationsMocks.notifyViaTerminal).toHaveBeenCalled(), - ); + expect(terminalNotificationsMocks.notifyViaTerminal).toHaveBeenCalled(); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('sends a macOS notification when a response completes while unfocused', async () => { @@ -732,35 +696,24 @@ describe('AppContainer State Management', () => { streamingState: currentStreamingState, })); - let unmount: (() => void) | undefined; - let rerender: ((tree: ReactElement) => void) | undefined; - - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - rerender = rendered.rerender; - }); + const { unmount, rerender } = await act(async () => renderAppContainer()); currentStreamingState = 'idle'; await act(async () => { - rerender?.(getAppContainer()); + rerender(getAppContainer()); }); - await waitFor(() => - expect( - terminalNotificationsMocks.buildRunEventNotificationContent, - ).toHaveBeenCalledWith( - expect.objectContaining({ - type: 'session_complete', - detail: 'Gemini CLI finished responding.', - }), - ), + expect( + terminalNotificationsMocks.buildRunEventNotificationContent, + ).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'session_complete', + detail: 'Gemini CLI finished responding.', + }), ); expect(terminalNotificationsMocks.notifyViaTerminal).toHaveBeenCalled(); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('sends completion notification when focus reporting is unavailable', async () => { @@ -774,34 +727,23 @@ describe('AppContainer State Management', () => { streamingState: currentStreamingState, })); - let unmount: (() => void) | undefined; - let rerender: ((tree: ReactElement) => void) | undefined; - - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - rerender = rendered.rerender; - }); + const { unmount, rerender } = await act(async () => renderAppContainer()); currentStreamingState = 'idle'; await act(async () => { - rerender?.(getAppContainer()); + rerender(getAppContainer()); }); - await waitFor(() => - expect( - terminalNotificationsMocks.buildRunEventNotificationContent, - ).toHaveBeenCalledWith( - expect.objectContaining({ - type: 'session_complete', - detail: 'Gemini CLI finished responding.', - }), - ), + expect( + terminalNotificationsMocks.buildRunEventNotificationContent, + ).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'session_complete', + detail: 'Gemini CLI finished responding.', + }), ); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('does not send completion notification when another action-required dialog is pending', async () => { @@ -819,27 +761,18 @@ describe('AppContainer State Management', () => { streamingState: currentStreamingState, })); - let unmount: (() => void) | undefined; - let rerender: ((tree: ReactElement) => void) | undefined; - - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - rerender = rendered.rerender; - }); + const { unmount, rerender } = await act(async () => renderAppContainer()); currentStreamingState = 'idle'; await act(async () => { - rerender?.(getAppContainer()); + rerender(getAppContainer()); }); expect( terminalNotificationsMocks.notifyViaTerminal, ).not.toHaveBeenCalled(); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('can send repeated attention notifications for the same key after pending state clears', async () => { @@ -875,24 +808,15 @@ describe('AppContainer State Management', () => { pendingHistoryItems, })); - let unmount: (() => void) | undefined; - let rerender: ((tree: ReactElement) => void) | undefined; + const { unmount, rerender } = await act(async () => renderAppContainer()); - await act(async () => { - const rendered = renderAppContainer(); - unmount = rendered.unmount; - rerender = rendered.rerender; - }); - - await waitFor(() => - expect( - terminalNotificationsMocks.notifyViaTerminal, - ).toHaveBeenCalledTimes(1), - ); + expect( + terminalNotificationsMocks.notifyViaTerminal, + ).toHaveBeenCalledTimes(1); pendingHistoryItems = []; await act(async () => { - rerender?.(getAppContainer()); + rerender(getAppContainer()); }); pendingHistoryItems = [ @@ -917,18 +841,14 @@ describe('AppContainer State Management', () => { }, ]; await act(async () => { - rerender?.(getAppContainer()); + rerender(getAppContainer()); }); - await waitFor(() => - expect( - terminalNotificationsMocks.notifyViaTerminal, - ).toHaveBeenCalledTimes(2), - ); + expect( + terminalNotificationsMocks.notifyViaTerminal, + ).toHaveBeenCalledTimes(2); - await act(async () => { - unmount?.(); - }); + unmount(); }); it('initializes with theme error from initialization result', async () => { @@ -937,68 +857,53 @@ describe('AppContainer State Management', () => { themeError: 'Failed to load theme', }; - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ initResult: initResultWithError, - }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }); - it('handles debug mode state', () => { + it('handles debug mode state', async () => { const debugConfig = makeFakeConfig(); vi.spyOn(debugConfig, 'getDebugMode').mockReturnValue(true); - expect(() => { - renderAppContainer({ config: debugConfig }); - }).not.toThrow(); + const { unmount } = await act(async () => + renderAppContainer({ config: debugConfig }), + ); + unmount(); }); }); describe('Context Providers', () => { it('provides AppContext with correct values', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ version: '2.0.0' }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => + renderAppContainer({ version: '2.0.0' }), + ); + expect(capturedUIState).toBeTruthy(); // Should render and unmount cleanly - expect(() => unmount!()).not.toThrow(); + unmount(); }); it('provides UIStateContext with state management', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + expect(capturedUIState).toBeTruthy(); + unmount(); }); it('provides UIActionsContext with action handlers', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + expect(capturedUIState).toBeTruthy(); + unmount(); }); it('provides ConfigContext with config object', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + expect(capturedUIState).toBeTruthy(); + unmount(); }); }); @@ -1011,13 +916,11 @@ describe('AppContainer State Management', () => { showMemoryUsage: false, }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ settings: settingsAllHidden }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => + renderAppContainer({ settings: settingsAllHidden }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }); it('handles settings with memory usage enabled', async () => { @@ -1025,13 +928,11 @@ describe('AppContainer State Management', () => { showMemoryUsage: true, }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ settings: settingsWithMemory }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => + renderAppContainer({ settings: settingsWithMemory }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }); }); @@ -1039,13 +940,11 @@ describe('AppContainer State Management', () => { it.each(['1.0.0', '2.1.3-beta', '3.0.0-nightly'])( 'handles version format: %s', async (version) => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ version }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => + renderAppContainer({ version }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }, ); }); @@ -1058,30 +957,30 @@ describe('AppContainer State Management', () => { }); // Should still render without crashing - errors should be handled internally - const { unmount } = renderAppContainer({ config: errorConfig }); + const { unmount } = await act(async () => + renderAppContainer({ config: errorConfig }), + ); unmount(); }); it('handles undefined settings gracefully', async () => { const undefinedSettings = createMockSettings(); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ settings: undefinedSettings }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - unmount!(); + const { unmount } = await act(async () => + renderAppContainer({ settings: undefinedSettings }), + ); + expect(capturedUIState).toBeTruthy(); + unmount(); }); }); describe('Provider Hierarchy', () => { - it('establishes correct provider nesting order', () => { + it('establishes correct provider nesting order', async () => { // This tests that all the context providers are properly nested // and that the component tree can be built without circular dependencies - const { unmount } = renderAppContainer(); + const { unmount } = await act(async () => renderAppContainer()); - expect(() => unmount()).not.toThrow(); + unmount(); }); }); @@ -1113,40 +1012,32 @@ describe('AppContainer State Management', () => { filePath: '/tmp/test-session.json', }; - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ config: mockConfig, settings: mockSettings, version: '1.0.0', initResult: mockInitResult, resumedSessionData: mockResumedSessionData, - }); - unmount = result.unmount; - }); - await act(async () => { - unmount(); - }); + }), + ); + unmount(); }); it('renders without resumed session data', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ config: mockConfig, settings: mockSettings, version: '1.0.0', initResult: mockInitResult, resumedSessionData: undefined, - }); - unmount = result.unmount; - }); - await act(async () => { - unmount(); - }); + }), + ); + unmount(); }); - it('initializes chat recording service when config has it', () => { + it('initializes chat recording service when config has it', async () => { const mockChatRecordingService = { initialize: vi.fn(), recordMessage: vi.fn(), @@ -1166,18 +1057,19 @@ describe('AppContainer State Management', () => { mockGeminiClient as unknown as ReturnType, ); - expect(() => { + const { unmount } = await act(async () => renderAppContainer({ config: configWithRecording, settings: mockSettings, version: '1.0.0', initResult: mockInitResult, - }); - }).not.toThrow(); + }), + ); + unmount(); }); }); describe('Session Recording Integration', () => { - it('provides chat recording service configuration', () => { + it('provides chat recording service configuration', async () => { const mockChatRecordingService = { initialize: vi.fn(), recordMessage: vi.fn(), @@ -1203,23 +1095,24 @@ describe('AppContainer State Management', () => { 'test-session-123', ); - expect(() => { + const { unmount } = await act(async () => renderAppContainer({ config: configWithRecording, settings: mockSettings, version: '1.0.0', initResult: mockInitResult, - }); - }).not.toThrow(); + }), + ); // Verify the recording service structure is correct expect(configWithRecording.getGeminiClient).toBeDefined(); expect(mockGeminiClient.getChatRecordingService).toBeDefined(); expect(mockChatRecordingService.initialize).toBeDefined(); expect(mockChatRecordingService.recordMessage).toBeDefined(); + unmount(); }); - it('handles session recording when messages are added', () => { + it('handles session recording when messages are added', async () => { const mockRecordMessage = vi.fn(); const mockRecordMessageTokens = vi.fn(); @@ -1242,22 +1135,25 @@ describe('AppContainer State Management', () => { mockGeminiClient as unknown as ReturnType, ); - renderAppContainer({ - config: configWithRecording, - settings: mockSettings, - version: '1.0.0', - initResult: mockInitResult, - }); + const { unmount } = await act(async () => + renderAppContainer({ + config: configWithRecording, + settings: mockSettings, + version: '1.0.0', + initResult: mockInitResult, + }), + ); // The actual recording happens through the useHistory hook // which would be triggered by user interactions expect(mockChatRecordingService.initialize).toBeDefined(); expect(mockChatRecordingService.recordMessage).toBeDefined(); + unmount(); }); }); describe('Session Resume Flow', () => { - it('accepts resumed session data', () => { + it('accepts resumed session data', async () => { const mockResumeChat = vi.fn(); const mockGeminiClient = { isInitialized: vi.fn(() => true), @@ -1303,22 +1199,23 @@ describe('AppContainer State Management', () => { filePath: '/tmp/resumed-session.json', }; - expect(() => { + const { unmount } = await act(async () => renderAppContainer({ config: configWithClient, settings: mockSettings, version: '1.0.0', initResult: mockInitResult, resumedSessionData: resumedData, - }); - }).not.toThrow(); + }), + ); // Verify the resume functionality structure is in place expect(mockGeminiClient.resumeChat).toBeDefined(); expect(resumedData.conversation.messages).toHaveLength(2); + unmount(); }); - it('does not attempt resume when client is not initialized', () => { + it('does not attempt resume when client is not initialized', async () => { const mockResumeChat = vi.fn(); const mockGeminiClient = { isInitialized: vi.fn(() => false), // Not initialized @@ -1343,21 +1240,24 @@ describe('AppContainer State Management', () => { filePath: '/tmp/session.json', }; - renderAppContainer({ - config: configWithClient, - settings: mockSettings, - version: '1.0.0', - initResult: mockInitResult, - resumedSessionData: resumedData, - }); + const { unmount } = await act(async () => + renderAppContainer({ + config: configWithClient, + settings: mockSettings, + version: '1.0.0', + initResult: mockInitResult, + resumedSessionData: resumedData, + }), + ); // Should not call resumeChat when client is not initialized expect(mockResumeChat).not.toHaveBeenCalled(); + unmount(); }); }); describe('Token Counting from Session Stats', () => { - it('tracks token counts from session messages', () => { + it('tracks token counts from session messages', async () => { // Session stats are provided through the SessionStatsProvider context // in the real app, not through the config directly const mockChatRecordingService = { @@ -1385,33 +1285,30 @@ describe('AppContainer State Management', () => { mockGeminiClient as unknown as ReturnType, ); - renderAppContainer({ - config: configWithRecording, - settings: mockSettings, - version: '1.0.0', - initResult: mockInitResult, - }); + const { unmount } = await act(async () => + renderAppContainer({ + config: configWithRecording, + settings: mockSettings, + version: '1.0.0', + initResult: mockInitResult, + }), + ); // In the actual app, these stats would be displayed in components // and updated as messages are processed through the recording service expect(mockChatRecordingService.recordMessageTokens).toBeDefined(); expect(mockChatRecordingService.getCurrentConversation).toBeDefined(); + unmount(); }); }); describe('Quota and Fallback Integration', () => { it('passes a null proQuotaRequest to UIStateContext by default', async () => { // The default mock from beforeEach already sets proQuotaRequest to null - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => { - // Assert that the context value is as expected - expect(capturedUIState.quota.proQuotaRequest).toBeNull(); - }); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + // Assert that the context value is as expected + expect(capturedUIState.quota.proQuotaRequest).toBeNull(); + unmount(); }); it('passes a valid proQuotaRequest to UIStateContext when provided by the hook', async () => { @@ -1427,16 +1324,10 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => { - // Assert: The mock request is correctly passed through the context - expect(capturedUIState.quota.proQuotaRequest).toEqual(mockRequest); - }); - unmount!(); + const { unmount } = await act(async () => renderAppContainer()); + // Assert: The mock request is correctly passed through the context + expect(capturedUIState.quota.proQuotaRequest).toEqual(mockRequest); + unmount(); }); it('passes the handleProQuotaChoice function to UIActionsContext', async () => { @@ -1448,22 +1339,16 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => { - // Assert: The action in the context is the mock handler we provided - expect(capturedUIActions.handleProQuotaChoice).toBe(mockHandler); - }); + const { unmount } = await act(async () => renderAppContainer()); + // Assert: The action in the context is the mock handler we provided + expect(capturedUIActions.handleProQuotaChoice).toBe(mockHandler); // You can even verify that the plumbed function is callable act(() => { capturedUIActions.handleProQuotaChoice('retry_later'); }); expect(mockHandler).toHaveBeenCalledWith('retry_later'); - unmount!(); + unmount(); }); }); @@ -1479,7 +1364,7 @@ describe('AppContainer State Management', () => { expect(stdout).toBe(mocks.mockStdout); }); - it('should update terminal title with Working… when showStatusInTitle is false', () => { + it('should update terminal title with Working… when showStatusInTitle is false', async () => { // Arrange: Set up mock settings with showStatusInTitle disabled const mockSettingsWithShowStatusFalse = createMockSettings({ ui: { @@ -1496,9 +1381,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithShowStatusFalse, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithShowStatusFalse, + }), + ); // Assert: Check that title was updated with "Working…" const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1512,7 +1399,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should use legacy terminal title when dynamicWindowTitle is false', () => { + it('should use legacy terminal title when dynamicWindowTitle is false', async () => { // Arrange: Set up mock settings with dynamicWindowTitle disabled const mockSettingsWithDynamicTitleFalse = createMockSettings({ ui: { @@ -1529,9 +1416,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithDynamicTitleFalse, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithDynamicTitleFalse, + }), + ); // Assert: Check that legacy title was used const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1545,7 +1434,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should not update terminal title when hideWindowTitle is true', () => { + it('should not update terminal title when hideWindowTitle is true', async () => { // Arrange: Set up mock settings with hideWindowTitle enabled const mockSettingsWithHideTitleTrue = createMockSettings({ ui: { @@ -1555,9 +1444,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithHideTitleTrue, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithHideTitleTrue, + }), + ); // Assert: Check that no title-related writes occurred const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1568,7 +1459,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should update terminal title with thought subject when in active state', () => { + it('should update terminal title with thought subject when in active state', async () => { // Arrange: Set up mock settings with showStatusInTitle enabled const mockSettingsWithTitleEnabled = createMockSettings({ ui: { @@ -1586,9 +1477,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Assert: Check that title was updated with thought subject and suffix const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1602,7 +1495,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should update terminal title with default text when in Idle state and no thought subject', () => { + it('should update terminal title with default text when in Idle state and no thought subject', async () => { // Arrange: Set up mock settings with showStatusInTitle enabled const mockSettingsWithTitleEnabled = createMockSettings({ ui: { @@ -1615,9 +1508,11 @@ describe('AppContainer State Management', () => { mockedUseGeminiStream.mockReturnValue(DEFAULT_GEMINI_STREAM_MOCK); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Assert: Check that title was updated with default Idle text const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1649,13 +1544,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ settings: mockSettingsWithTitleEnabled, - }); - unmount = result.unmount; - }); + }), + ); // Assert: Check that title was updated with confirmation text const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1666,7 +1559,7 @@ describe('AppContainer State Management', () => { expect(titleWrites[0][0]).toBe( `\x1b]0;${'✋ Action Required (workspace)'.padEnd(80, ' ')}\x07`, ); - unmount!(); + unmount(); }); describe('Shell Focus Action Required', () => { @@ -1712,9 +1605,11 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'isInteractiveShellEnabled').mockReturnValue(true); // Act: Render the container (embeddedShellFocused is false by default in state) - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Initially it should show the working status const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1773,9 +1668,11 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'isInteractive').mockReturnValue(true); vi.spyOn(mockConfig, 'isInteractiveShellEnabled').mockReturnValue(true); - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Fast-forward time by 65 seconds - should still NOT be Action Required await act(async () => { @@ -1830,9 +1727,11 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'isInteractive').mockReturnValue(true); vi.spyOn(mockConfig, 'isInteractiveShellEnabled').mockReturnValue(true); - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Fast-forward time by 65 seconds await act(async () => { @@ -1875,9 +1774,11 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'isInteractiveShellEnabled').mockReturnValue(true); // Act: Render the container - const { unmount, rerender } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount, rerender } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Fast-forward time by 20 seconds await act(async () => { @@ -1931,7 +1832,7 @@ describe('AppContainer State Management', () => { }); }); - it('should pad title to exactly 80 characters', () => { + it('should pad title to exactly 80 characters', async () => { // Arrange: Set up mock settings with showStatusInTitle enabled const mockSettingsWithTitleEnabled = createMockSettings({ ui: { @@ -1949,9 +1850,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Assert: Check that title is padded to exactly 80 characters const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1966,7 +1869,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should use correct ANSI escape code format', () => { + it('should use correct ANSI escape code format', async () => { // Arrange: Set up mock settings with showStatusInTitle enabled const mockSettingsWithTitleEnabled = createMockSettings({ ui: { @@ -1984,9 +1887,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleEnabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleEnabled, + }), + ); // Assert: Check that the correct ANSI escape sequence is used const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -1999,7 +1904,7 @@ describe('AppContainer State Management', () => { unmount(); }); - it('should use CLI_TITLE environment variable when set', () => { + it('should use CLI_TITLE environment variable when set', async () => { // Arrange: Set up mock settings with showStatusInTitle disabled (so it shows suffix) const mockSettingsWithTitleDisabled = createMockSettings({ ui: { @@ -2018,9 +1923,11 @@ describe('AppContainer State Management', () => { }); // Act: Render the container - const { unmount } = renderAppContainer({ - settings: mockSettingsWithTitleDisabled, - }); + const { unmount } = await act(async () => + renderAppContainer({ + settings: mockSettingsWithTitleDisabled, + }), + ); // Assert: Check that title was updated with CLI_TITLE value const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => @@ -2046,7 +1953,7 @@ describe('AppContainer State Management', () => { }); it('should set and clear the queue error message after a timeout', async () => { - const { rerender, unmount } = renderAppContainer(); + const { rerender, unmount } = await act(async () => renderAppContainer()); await act(async () => { vi.advanceTimersByTime(0); }); @@ -2068,7 +1975,7 @@ describe('AppContainer State Management', () => { }); it('should reset the timer if a new error message is set', async () => { - const { rerender, unmount } = renderAppContainer(); + const { rerender, unmount } = await act(async () => renderAppContainer()); await act(async () => { vi.advanceTimersByTime(0); }); @@ -2110,11 +2017,11 @@ describe('AppContainer State Management', () => { let mockCancelOngoingRequest: Mock; let rerender: () => void; let unmount: () => void; - let stdin: ReturnType['stdin']; + let stdin: Awaited>['stdin']; // Helper function to reduce boilerplate in tests const setupKeypressTest = async () => { - const renderResult = renderAppContainer(); + const renderResult = await act(async () => renderAppContainer()); stdin = renderResult.stdin; await act(async () => { vi.advanceTimersByTime(0); @@ -2328,7 +2235,7 @@ describe('AppContainer State Management', () => { activePtyId: 1, }); - const renderResult = render(getAppContainer()); + const renderResult = await act(async () => render(getAppContainer())); await act(async () => { vi.advanceTimersByTime(0); }); @@ -2446,7 +2353,7 @@ describe('AppContainer State Management', () => { let unmount: () => void; const setupShortcutsVisibilityTest = async () => { - const renderResult = renderAppContainer(); + const renderResult = await act(async () => renderAppContainer()); await act(async () => { vi.advanceTimersByTime(0); }); @@ -2522,9 +2429,7 @@ describe('AppContainer State Management', () => { await act(async () => { rerender(); }); - await waitFor(() => { - expect(capturedUIState.shortcutsHelpVisible).toBe(false); - }); + expect(capturedUIState.shortcutsHelpVisible).toBe(false); unmount(); }); @@ -2553,9 +2458,7 @@ describe('AppContainer State Management', () => { await act(async () => { rerender(); }); - await waitFor(() => { - expect(capturedUIState.shortcutsHelpVisible).toBe(false); - }); + expect(capturedUIState.shortcutsHelpVisible).toBe(false); unmount(); }); @@ -2564,7 +2467,7 @@ describe('AppContainer State Management', () => { describe('Copy Mode (CTRL+S)', () => { let rerender: () => void; let unmount: () => void; - let stdin: ReturnType['stdin']; + let stdin: Awaited>['stdin']; const setupCopyModeTest = async ( isAlternateMode = false, @@ -2602,7 +2505,7 @@ describe('AppContainer State Management', () => { ); - const renderResult = render(getTree(testSettings)); + const renderResult = await act(async () => render(getTree(testSettings))); stdin = renderResult.stdin; await act(async () => { vi.advanceTimersByTime(0); @@ -2792,15 +2695,10 @@ describe('AppContainer State Management', () => { closeModelDialog: vi.fn(), }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); expect(capturedUIState.isModelDialogOpen).toBe(true); - unmount!(); + unmount(); }); it('should provide model dialog actions in the UIActionsContext', async () => { @@ -2812,45 +2710,29 @@ describe('AppContainer State Management', () => { closeModelDialog: mockCloseModelDialog, }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); // Verify that the actions are correctly passed through context act(() => { capturedUIActions.closeModelDialog(); }); expect(mockCloseModelDialog).toHaveBeenCalled(); - unmount!(); + unmount(); }); }); describe('Agent Configuration Dialog Integration', () => { it('should initialize with dialog closed and no agent selected', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); - + const { unmount } = await act(async () => renderAppContainer()); expect(capturedUIState.isAgentConfigDialogOpen).toBe(false); expect(capturedUIState.selectedAgentName).toBeUndefined(); expect(capturedUIState.selectedAgentDisplayName).toBeUndefined(); expect(capturedUIState.selectedAgentDefinition).toBeUndefined(); - unmount!(); + unmount(); }); it('should update state when openAgentConfigDialog is called', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); const agentDefinition = { name: 'test-agent' }; act(() => { @@ -2865,16 +2747,11 @@ describe('AppContainer State Management', () => { expect(capturedUIState.selectedAgentName).toBe('test-agent'); expect(capturedUIState.selectedAgentDisplayName).toBe('Test Agent'); expect(capturedUIState.selectedAgentDefinition).toEqual(agentDefinition); - unmount!(); + unmount(); }); it('should clear state when closeAgentConfigDialog is called', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); const agentDefinition = { name: 'test-agent' }; act(() => { @@ -2895,31 +2772,26 @@ describe('AppContainer State Management', () => { expect(capturedUIState.selectedAgentName).toBeUndefined(); expect(capturedUIState.selectedAgentDisplayName).toBeUndefined(); expect(capturedUIState.selectedAgentDefinition).toBeUndefined(); - unmount!(); + unmount(); }); }); describe('CoreEvents Integration', () => { it('subscribes to UserFeedback and drains backlog on mount', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); expect(mockCoreEvents.on).toHaveBeenCalledWith( CoreEvent.UserFeedback, expect.any(Function), ); expect(mockCoreEvents.drainBacklogs).toHaveBeenCalledTimes(1); - unmount!(); + unmount(); }); it('unsubscribes from UserFeedback on unmount', async () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => expect(capturedUIState).toBeTruthy()); @@ -2935,7 +2807,7 @@ describe('AppContainer State Management', () => { it('adds history item when UserFeedback event is received', async () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => expect(capturedUIState).toBeTruthy()); @@ -2971,7 +2843,7 @@ describe('AppContainer State Management', () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => { @@ -3004,7 +2876,7 @@ describe('AppContainer State Management', () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => expect(capturedUIState).toBeTruthy()); @@ -3016,7 +2888,7 @@ describe('AppContainer State Management', () => { it('handles consent request events', async () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => expect(capturedUIState).toBeTruthy()); @@ -3053,7 +2925,7 @@ describe('AppContainer State Management', () => { it('unsubscribes from ConsentRequest on unmount', async () => { let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => expect(capturedUIState).toBeTruthy()); @@ -3076,7 +2948,7 @@ describe('AppContainer State Management', () => { }); let unmount: () => void; await act(async () => { - const result = renderAppContainer(); + const result = await renderAppContainer(); unmount = result.unmount; }); await waitFor(() => { @@ -3104,12 +2976,7 @@ describe('AppContainer State Management', () => { }); it('preserves buffer when cancelling, even if empty (user is in control)', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); const { onCancelSubmit } = extractUseGeminiStreamArgs( mockedUseGeminiStream.mock.lastCall!, @@ -3122,7 +2989,7 @@ describe('AppContainer State Management', () => { // Should NOT modify buffer when cancelling - user is in control expect(mockSetText).not.toHaveBeenCalled(); - unmount!(); + unmount(); }); it('preserves prompt text when cancelling streaming, even if same as last message (regression test for issue #13387)', async () => { @@ -3140,12 +3007,7 @@ describe('AppContainer State Management', () => { initializeFromLogger: vi.fn(), }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); const { onCancelSubmit } = extractUseGeminiStreamArgs( mockedUseGeminiStream.mock.lastCall!, @@ -3159,7 +3021,7 @@ describe('AppContainer State Management', () => { // Should NOT call setText - prompt should be preserved regardless of content expect(mockSetText).not.toHaveBeenCalled(); - unmount!(); + unmount(); }); it('restores the prompt when onCancelSubmit is called with shouldRestorePrompt=true (or undefined)', async () => { @@ -3170,14 +3032,8 @@ describe('AppContainer State Management', () => { initializeFromLogger: vi.fn(), }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => - expect(capturedUIState.userMessages).toContain('previous message'), - ); + const { unmount } = await act(async () => renderAppContainer()); + expect(capturedUIState.userMessages).toContain('previous message'); const { onCancelSubmit } = extractUseGeminiStreamArgs( mockedUseGeminiStream.mock.lastCall!, @@ -3187,11 +3043,9 @@ describe('AppContainer State Management', () => { onCancelSubmit(true); }); - await waitFor(() => { - expect(mockSetText).toHaveBeenCalledWith('previous message'); - }); + expect(mockSetText).toHaveBeenCalledWith('previous message'); - unmount!(); + unmount(); }); it('input history is independent from conversation history (survives /clear)', async () => { @@ -3204,18 +3058,10 @@ describe('AppContainer State Management', () => { initializeFromLogger: vi.fn(), }); - let rerender: (tree: ReactElement) => void; - let unmount; - await act(async () => { - const result = renderAppContainer(); - rerender = result.rerender; - unmount = result.unmount; - }); + const { rerender, unmount } = await act(async () => renderAppContainer()); // Verify userMessages is populated from inputHistory - await waitFor(() => - expect(capturedUIState.userMessages).toContain('first prompt'), - ); + expect(capturedUIState.userMessages).toContain('first prompt'); expect(capturedUIState.userMessages).toContain('second prompt'); // Clear the conversation history (simulating /clear command) @@ -3238,7 +3084,7 @@ describe('AppContainer State Management', () => { expect(capturedUIState.userMessages).toContain('first prompt'); expect(capturedUIState.userMessages).toContain('second prompt'); - unmount!(); + unmount(); }); }); @@ -3253,14 +3099,10 @@ describe('AppContainer State Management', () => { // Clear previous calls mocks.mockStdout.write.mockClear(); - let compUnmount: () => void = () => {}; - await act(async () => { - const { unmount } = renderAppContainer(); - compUnmount = unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); // Allow async effects to run - await waitFor(() => expect(capturedUIState).toBeTruthy()); + expect(capturedUIState).toBeTruthy(); // Wait for fetchBannerTexts to complete await act(async () => { @@ -3273,7 +3115,7 @@ describe('AppContainer State Management', () => { ); expect(clearTerminalCalls).toHaveLength(0); - compUnmount(); + unmount(); }); }); @@ -3284,14 +3126,13 @@ describe('AppContainer State Management', () => { ); vi.mocked(checkPermissions).mockResolvedValue([]); - let unmount: () => void; - await act(async () => { - unmount = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ settings: createMockSettings({ ui: { useAlternateBuffer: false } }), - }).unmount; - }); + }), + ); - await waitFor(() => expect(capturedUIActions).toBeTruthy()); + expect(capturedUIActions).toBeTruthy(); // Expand first act(() => capturedUIActions.setConstrainHeight(false)); @@ -3309,7 +3150,7 @@ describe('AppContainer State Management', () => { expect(mocks.mockStdout.write).toHaveBeenCalledWith( ansiEscapes.clearTerminal, ); - unmount!(); + unmount(); }); it('resets expansion state on submission when in alternate buffer without clearing terminal', async () => { @@ -3320,14 +3161,13 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'getUseAlternateBuffer').mockReturnValue(true); - let unmount: () => void; - await act(async () => { - unmount = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ settings: createMockSettings({ ui: { useAlternateBuffer: true } }), - }).unmount; - }); + }), + ); - await waitFor(() => expect(capturedUIActions).toBeTruthy()); + expect(capturedUIActions).toBeTruthy(); // Expand first act(() => capturedUIActions.setConstrainHeight(false)); @@ -3345,7 +3185,7 @@ describe('AppContainer State Management', () => { expect(mocks.mockStdout.write).not.toHaveBeenCalledWith( ansiEscapes.clearTerminal, ); - unmount!(); + unmount(); }); }); @@ -3358,13 +3198,9 @@ describe('AppContainer State Management', () => { vi.useRealTimers(); }); - it('sets showIsExpandableHint when overflow occurs in Standard Mode and hides after 10s', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + it('should set showIsExpandableHint when overflow occurs in Standard Mode and hides after 10s', async () => { + const { unmount } = await act(async () => renderAppContainer()); + await waitFor(() => expect(capturedOverflowActions).toBeTruthy()); // Trigger overflow act(() => { @@ -3390,16 +3226,12 @@ describe('AppContainer State Management', () => { expect(capturedUIState.showIsExpandableHint).toBe(false); }); - unmount!(); + unmount(); }); it('resets the hint timer when a new component overflows (overflowingIdsSize increases)', async () => { - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { unmount } = await act(async () => renderAppContainer()); + await waitFor(() => expect(capturedOverflowActions).toBeTruthy()); // 1. Trigger first overflow act(() => { @@ -3447,18 +3279,12 @@ describe('AppContainer State Management', () => { expect(capturedUIState.showIsExpandableHint).toBe(false); }); - unmount!(); + unmount(); }); it('toggles expansion state and resets the hint timer when Ctrl+O is pressed in Standard Mode', async () => { - let unmount: () => void; - let stdin: ReturnType['stdin']; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - stdin = result.stdin; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { stdin, unmount } = await act(async () => renderAppContainer()); + await waitFor(() => expect(capturedOverflowActions).toBeTruthy()); // Initial state is constrainHeight = true expect(capturedUIState.constrainHeight).toBe(true); @@ -3483,10 +3309,8 @@ describe('AppContainer State Management', () => { stdin.write('\x0f'); // \x0f is Ctrl+O }); - await waitFor(() => { - // constrainHeight should toggle - expect(capturedUIState.constrainHeight).toBe(false); - }); + // constrainHeight should toggle + expect(capturedUIState.constrainHeight).toBe(false); // Advance enough that the original timer would have expired if it hadn't reset act(() => { @@ -3505,18 +3329,12 @@ describe('AppContainer State Management', () => { expect(capturedUIState.showIsExpandableHint).toBe(false); }); - unmount!(); + unmount(); }); it('toggles Ctrl+O multiple times and verifies the hint disappears exactly after the last toggle', async () => { - let unmount: () => void; - let stdin: ReturnType['stdin']; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - stdin = result.stdin; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + const { stdin, unmount } = await act(async () => renderAppContainer()); + await waitFor(() => expect(capturedOverflowActions).toBeTruthy()); // Initial state is constrainHeight = true expect(capturedUIState.constrainHeight).toBe(true); @@ -3540,9 +3358,7 @@ describe('AppContainer State Management', () => { act(() => { stdin.write('\x0f'); // Ctrl+O }); - await waitFor(() => { - expect(capturedUIState.constrainHeight).toBe(false); - }); + expect(capturedUIState.constrainHeight).toBe(false); // Wait 1 second act(() => { @@ -3554,9 +3370,7 @@ describe('AppContainer State Management', () => { act(() => { stdin.write('\x0f'); // Ctrl+O }); - await waitFor(() => { - expect(capturedUIState.constrainHeight).toBe(true); - }); + expect(capturedUIState.constrainHeight).toBe(true); // Wait 1 second act(() => { @@ -3568,9 +3382,7 @@ describe('AppContainer State Management', () => { act(() => { stdin.write('\x0f'); // Ctrl+O }); - await waitFor(() => { - expect(capturedUIState.constrainHeight).toBe(false); - }); + expect(capturedUIState.constrainHeight).toBe(false); // Now we wait just before the timeout from the LAST toggle. // It should still be true. @@ -3588,7 +3400,7 @@ describe('AppContainer State Management', () => { expect(capturedUIState.showIsExpandableHint).toBe(false); }); - unmount!(); + unmount(); }); it('DOES set showIsExpandableHint when overflow occurs in Alternate Buffer Mode', async () => { @@ -3598,14 +3410,12 @@ describe('AppContainer State Management', () => { vi.spyOn(mockConfig, 'getUseAlternateBuffer').mockReturnValue(true); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer({ + const { unmount } = await act(async () => + renderAppContainer({ settings: settingsWithAlternateBuffer, - }); - unmount = result.unmount; - }); - await waitFor(() => expect(capturedUIState).toBeTruthy()); + }), + ); + await waitFor(() => expect(capturedOverflowActions).toBeTruthy()); // Trigger overflow act(() => { @@ -3617,7 +3427,7 @@ describe('AppContainer State Management', () => { expect(capturedUIState.showIsExpandableHint).toBe(true); }); - unmount!(); + unmount(); }); }); @@ -3628,10 +3438,9 @@ describe('AppContainer State Management', () => { ); vi.mocked(checkPermissions).mockResolvedValue(['/test/file.txt']); - let unmount: () => void; - await act(async () => (unmount = renderAppContainer().unmount)); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => expect(capturedUIActions).toBeTruthy()); + expect(capturedUIActions).toBeTruthy(); await act(async () => capturedUIActions.handleFinalSubmit('read @file.txt'), @@ -3641,7 +3450,7 @@ describe('AppContainer State Management', () => { expect(capturedUIState.permissionConfirmationRequest?.files).toEqual([ '/test/file.txt', ]); - await act(async () => unmount!()); + unmount(); }); it.each([true, false])( @@ -3657,10 +3466,9 @@ describe('AppContainer State Management', () => { ); const { submitQuery } = mockedUseGeminiStream(); - let unmount: () => void; - await act(async () => (unmount = renderAppContainer().unmount)); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => expect(capturedUIActions).toBeTruthy()); + expect(capturedUIActions).toBeTruthy(); await act(async () => capturedUIActions.handleFinalSubmit('read @file.txt'), @@ -3679,7 +3487,7 @@ describe('AppContainer State Management', () => { } expect(submitQuery).toHaveBeenCalledWith('read @file.txt'); expect(capturedUIState.permissionConfirmationRequest).toBeNull(); - await act(async () => unmount!()); + unmount(); }, ); }); @@ -3692,17 +3500,11 @@ describe('AppContainer State Management', () => { pendingHistoryItems: [], }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => { - expect(capturedUIState).toBeTruthy(); - expect(capturedUIState.allowPlanMode).toBe(true); - }); - unmount!(); + expect(capturedUIState).toBeTruthy(); + expect(capturedUIState.allowPlanMode).toBe(true); + unmount(); }); it('should NOT allow plan mode when disabled in config', async () => { @@ -3712,17 +3514,11 @@ describe('AppContainer State Management', () => { pendingHistoryItems: [], }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => { - expect(capturedUIState).toBeTruthy(); - expect(capturedUIState.allowPlanMode).toBe(false); - }); - unmount!(); + expect(capturedUIState).toBeTruthy(); + expect(capturedUIState.allowPlanMode).toBe(false); + unmount(); }); it('should NOT allow plan mode when streaming', async () => { @@ -3733,17 +3529,11 @@ describe('AppContainer State Management', () => { pendingHistoryItems: [], }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => { - expect(capturedUIState).toBeTruthy(); - expect(capturedUIState.allowPlanMode).toBe(false); - }); - unmount!(); + expect(capturedUIState).toBeTruthy(); + expect(capturedUIState.allowPlanMode).toBe(false); + unmount(); }); it('should NOT allow plan mode when a tool is awaiting confirmation', async () => { @@ -3764,17 +3554,11 @@ describe('AppContainer State Management', () => { ], }); - let unmount: () => void; - await act(async () => { - const result = renderAppContainer(); - unmount = result.unmount; - }); + const { unmount } = await act(async () => renderAppContainer()); - await waitFor(() => { - expect(capturedUIState).toBeTruthy(); - expect(capturedUIState.allowPlanMode).toBe(false); - }); - unmount!(); + expect(capturedUIState).toBeTruthy(); + expect(capturedUIState.allowPlanMode).toBe(false); + unmount(); }); }); }); diff --git a/packages/cli/src/ui/IdeIntegrationNudge.test.tsx b/packages/cli/src/ui/IdeIntegrationNudge.test.tsx index 5df3534f12..eb3e6a3e4c 100644 --- a/packages/cli/src/ui/IdeIntegrationNudge.test.tsx +++ b/packages/cli/src/ui/IdeIntegrationNudge.test.tsx @@ -53,10 +53,9 @@ describe('IdeIntegrationNudge', () => { }); it('renders correctly with default options', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const frame = lastFrame(); expect(frame).toContain('Do you want to connect VS Code to Gemini CLI?'); @@ -72,8 +71,6 @@ describe('IdeIntegrationNudge', () => { , ); - await waitUntilReady(); - // "Yes" is the first option and selected by default usually. await act(async () => { stdin.write('\r'); @@ -93,8 +90,6 @@ describe('IdeIntegrationNudge', () => { , ); - await waitUntilReady(); - // Navigate down to "No (esc)" await act(async () => { stdin.write('\u001B[B'); // Down arrow @@ -119,8 +114,6 @@ describe('IdeIntegrationNudge', () => { , ); - await waitUntilReady(); - // Navigate down to "No, don't ask again" await act(async () => { stdin.write('\u001B[B'); // Down arrow @@ -150,8 +143,6 @@ describe('IdeIntegrationNudge', () => { , ); - await waitUntilReady(); - // Press Escape await act(async () => { stdin.write('\u001B'); @@ -178,8 +169,6 @@ describe('IdeIntegrationNudge', () => { , ); - await waitUntilReady(); - const frame = lastFrame(); expect(frame).toContain( diff --git a/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx b/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx index b8de6adb0b..d46e0295a1 100644 --- a/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx +++ b/packages/cli/src/ui/auth/ApiAuthDialog.test.tsx @@ -73,23 +73,21 @@ describe('ApiAuthDialog', () => { }); it('renders correctly', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('renders with a defaultValue', async () => { - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); expect(mockedUseTextBuffer).toHaveBeenCalledWith( expect.objectContaining({ initialText: 'test-key', @@ -113,10 +111,9 @@ describe('ApiAuthDialog', () => { 'calls $expectedCall.name when $keyName is pressed', async ({ keyName, sequence, expectedCall, args }) => { mockBuffer.text = 'submitted-key'; // Set for the onSubmit case - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); // calls[0] is the ApiAuthDialog's useKeypress (Ctrl+C handler) // calls[1] is the TextInput's useKeypress (typing handler) const keypressHandler = mockedUseKeypress.mock.calls[1][0]; @@ -136,24 +133,22 @@ describe('ApiAuthDialog', () => { ); it('displays an error message', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Invalid API Key'); unmount(); }); it('calls clearApiKey and clears buffer when Ctrl+C is pressed', async () => { - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); // Call 0 is ApiAuthDialog (isActive: true) // Call 1 is TextInput (isActive: true, priority: true) const keypressHandler = mockedUseKeypress.mock.calls[0][0]; diff --git a/packages/cli/src/ui/auth/AuthDialog.test.tsx b/packages/cli/src/ui/auth/AuthDialog.test.tsx index 878b2a8ee0..4837a71490 100644 --- a/packages/cli/src/ui/auth/AuthDialog.test.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.test.tsx @@ -143,10 +143,9 @@ describe('AuthDialog', () => { for (const [key, value] of Object.entries(env)) { vi.stubEnv(key, value as string); } - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const items = mockedRadioButtonSelect.mock.calls[0][0].items; for (const item of shouldContain) { expect(items).toContainEqual(item); @@ -161,10 +160,7 @@ describe('AuthDialog', () => { it('filters auth types when enforcedType is set', async () => { props.settings.merged.security.auth.enforcedType = AuthType.USE_GEMINI; - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const items = mockedRadioButtonSelect.mock.calls[0][0].items; expect(items).toHaveLength(1); expect(items[0].value).toBe(AuthType.USE_GEMINI); @@ -173,10 +169,7 @@ describe('AuthDialog', () => { it('sets initial index to 0 when enforcedType is set', async () => { props.settings.merged.security.auth.enforcedType = AuthType.USE_GEMINI; - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { initialIndex } = mockedRadioButtonSelect.mock.calls[0][0]; expect(initialIndex).toBe(0); unmount(); @@ -213,10 +206,7 @@ describe('AuthDialog', () => { }, ])('selects initial auth type $desc', async ({ setup, expected }) => { setup(); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { items, initialIndex } = mockedRadioButtonSelect.mock.calls[0][0]; expect(items[initialIndex].value).toBe(expected); unmount(); @@ -226,10 +216,7 @@ describe('AuthDialog', () => { describe('handleAuthSelect', () => { it('calls onAuthError if validation fails', async () => { mockedValidateAuthMethod.mockReturnValue('Invalid method'); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; handleAuthSelect(AuthType.USE_GEMINI); @@ -245,10 +232,7 @@ describe('AuthDialog', () => { it('sets auth context with requiresRestart: true for LOGIN_WITH_GOOGLE', async () => { mockedValidateAuthMethod.mockReturnValue(null); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.LOGIN_WITH_GOOGLE); @@ -261,10 +245,7 @@ describe('AuthDialog', () => { it('sets auth context with empty object for other auth types', async () => { mockedValidateAuthMethod.mockReturnValue(null); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.USE_GEMINI); @@ -278,10 +259,7 @@ describe('AuthDialog', () => { vi.stubEnv('GEMINI_API_KEY', 'test-key-from-env'); // props.settings.merged.security.auth.selectedType is undefined here, simulating initial setup - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.USE_GEMINI); @@ -297,10 +275,7 @@ describe('AuthDialog', () => { vi.stubEnv('GEMINI_API_KEY', ''); // Empty string // props.settings.merged.security.auth.selectedType is undefined here - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.USE_GEMINI); @@ -316,10 +291,7 @@ describe('AuthDialog', () => { // process.env['GEMINI_API_KEY'] is not set // props.settings.merged.security.auth.selectedType is undefined here, simulating initial setup - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.USE_GEMINI); @@ -337,10 +309,7 @@ describe('AuthDialog', () => { props.settings.merged.security.auth.selectedType = AuthType.LOGIN_WITH_GOOGLE; - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await handleAuthSelect(AuthType.USE_GEMINI); @@ -360,10 +329,7 @@ describe('AuthDialog', () => { vi.mocked(props.config.isBrowserLaunchSuppressed).mockReturnValue(true); mockedValidateAuthMethod.mockReturnValue(null); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const { onSelect: handleAuthSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await act(async () => { @@ -383,10 +349,9 @@ describe('AuthDialog', () => { it('displays authError when provided', async () => { props.authError = 'Something went wrong'; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Something went wrong'); unmount(); }); @@ -429,10 +394,7 @@ describe('AuthDialog', () => { }, ])('$desc', async ({ setup, expectations }) => { setup(); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; keypressHandler({ name: 'escape' }); expectations(props); @@ -442,30 +404,27 @@ describe('AuthDialog', () => { describe('Snapshots', () => { it('renders correctly with default props', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('renders correctly with auth error', async () => { props.authError = 'Something went wrong'; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('renders correctly with enforced auth type', async () => { props.settings.merged.security.auth.enforcedType = AuthType.USE_GEMINI; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); diff --git a/packages/cli/src/ui/auth/AuthInProgress.test.tsx b/packages/cli/src/ui/auth/AuthInProgress.test.tsx index bd6a3cb126..1c392be28d 100644 --- a/packages/cli/src/ui/auth/AuthInProgress.test.tsx +++ b/packages/cli/src/ui/auth/AuthInProgress.test.tsx @@ -55,20 +55,18 @@ describe('AuthInProgress', () => { }); it('renders initial state with spinner', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toContain('[Spinner] Waiting for authentication...'); expect(lastFrame()).toContain('Press Esc or Ctrl+C to cancel'); unmount(); }); it('calls onTimeout when ESC is pressed', async () => { - const { waitUntilReady, unmount } = render( + const { waitUntilReady, unmount } = await render( , ); - await waitUntilReady(); const keypressHandler = vi.mocked(useKeypress).mock.calls[0][0]; await act(async () => { @@ -84,10 +82,9 @@ describe('AuthInProgress', () => { }); it('calls onTimeout when Ctrl+C is pressed', async () => { - const { waitUntilReady, unmount } = render( + const { waitUntilReady, unmount } = await render( , ); - await waitUntilReady(); const keypressHandler = vi.mocked(useKeypress).mock.calls[0][0]; await act(async () => { @@ -100,10 +97,9 @@ describe('AuthInProgress', () => { }); it('calls onTimeout and shows timeout message after 3 minutes', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, waitUntilReady, unmount } = await render( , ); - await waitUntilReady(); await act(async () => { vi.advanceTimersByTime(180000); @@ -116,10 +112,7 @@ describe('AuthInProgress', () => { }); it('clears timer on unmount', async () => { - const { waitUntilReady, unmount } = render( - , - ); - await waitUntilReady(); + const { unmount } = await render(); await act(async () => { unmount(); diff --git a/packages/cli/src/ui/auth/BannedAccountDialog.test.tsx b/packages/cli/src/ui/auth/BannedAccountDialog.test.tsx index 0670c81bc9..4b5d44e6d5 100644 --- a/packages/cli/src/ui/auth/BannedAccountDialog.test.tsx +++ b/packages/cli/src/ui/auth/BannedAccountDialog.test.tsx @@ -73,14 +73,13 @@ describe('BannedAccountDialog', () => { }); it('renders the suspension message from accountSuspensionInfo', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const frame = lastFrame(); expect(frame).toContain('Account Suspended'); expect(frame).toContain('violation of Terms of Service'); @@ -89,14 +88,13 @@ describe('BannedAccountDialog', () => { }); it('renders menu options with appeal link text from response', async () => { - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const items = mockedRadioButtonSelect.mock.calls[0][0].items; expect(items).toHaveLength(3); expect(items[0].label).toBe('Appeal Here'); @@ -109,14 +107,13 @@ describe('BannedAccountDialog', () => { const infoWithoutUrl: AccountSuspensionInfo = { message: 'Account suspended.', }; - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const items = mockedRadioButtonSelect.mock.calls[0][0].items; expect(items).toHaveLength(2); expect(items[0].label).toBe('Change authentication'); @@ -129,28 +126,26 @@ describe('BannedAccountDialog', () => { message: 'Account suspended.', appealUrl: 'https://example.com/appeal', }; - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const items = mockedRadioButtonSelect.mock.calls[0][0].items; expect(items[0].label).toBe('Open the Google Form'); unmount(); }); it('opens browser when appeal option is selected', async () => { - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const { onSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await onSelect('open_form'); expect(mockedOpenBrowser).toHaveBeenCalledWith( @@ -162,14 +157,13 @@ describe('BannedAccountDialog', () => { it('shows URL when browser cannot be launched', async () => { mockedShouldLaunchBrowser.mockReturnValue(false); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const { onSelect } = mockedRadioButtonSelect.mock.calls[0][0]; onSelect('open_form'); await waitFor(() => { @@ -180,14 +174,13 @@ describe('BannedAccountDialog', () => { }); it('calls onExit when "Exit" is selected', async () => { - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const { onSelect } = mockedRadioButtonSelect.mock.calls[0][0]; await onSelect('exit'); expect(mockedRunExitCleanup).toHaveBeenCalled(); @@ -196,14 +189,13 @@ describe('BannedAccountDialog', () => { }); it('calls onChangeAuth when "Change authentication" is selected', async () => { - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const { onSelect } = mockedRadioButtonSelect.mock.calls[0][0]; onSelect('change_auth'); expect(onChangeAuth).toHaveBeenCalled(); @@ -212,14 +204,13 @@ describe('BannedAccountDialog', () => { }); it('exits on escape key', async () => { - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; const result = keypressHandler({ name: 'escape' }); expect(result).toBe(true); @@ -227,14 +218,13 @@ describe('BannedAccountDialog', () => { }); it('renders snapshot correctly', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); diff --git a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx index 77310e3069..4dd13a3334 100644 --- a/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx +++ b/packages/cli/src/ui/auth/LoginWithGoogleRestartDialog.test.tsx @@ -45,25 +45,23 @@ describe('LoginWithGoogleRestartDialog', () => { }); it('renders correctly', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('calls onDismiss when escape is pressed', async () => { - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; keypressHandler({ @@ -83,13 +81,12 @@ describe('LoginWithGoogleRestartDialog', () => { async (keyName) => { vi.useFakeTimers(); - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); const keypressHandler = mockedUseKeypress.mock.calls[0][0]; keypressHandler({ diff --git a/packages/cli/src/ui/auth/useAuth.test.tsx b/packages/cli/src/ui/auth/useAuth.test.tsx index f236428ff1..8d51e46a64 100644 --- a/packages/cli/src/ui/auth/useAuth.test.tsx +++ b/packages/cli/src/ui/auth/useAuth.test.tsx @@ -4,15 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - describe, - it, - expect, - vi, - beforeEach, - afterEach, - type Mock, -} from 'vitest'; +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { act } from 'react'; import { renderHook } from '../../test-utils/render.js'; import { useAuthCommand, validateAuthMethodWithSettings } from './useAuth.js'; import { @@ -22,7 +15,6 @@ import { } from '@google/gemini-cli-core'; import { AuthState } from '../types.js'; import type { LoadedSettings } from '../../config/settings.js'; -import { waitFor } from '../../test-utils/async.js'; // Mock dependencies const mockLoadApiKey = vi.fn(); @@ -142,171 +134,202 @@ describe('useAuth', () => { }, }) as LoadedSettings; + let deferredRefreshAuth: { + resolve: () => void; + reject: (e: Error) => void; + }; + + beforeEach(() => { + vi.mocked(mockConfig.refreshAuth).mockImplementation( + () => + new Promise((resolve, reject) => { + deferredRefreshAuth = { resolve, reject }; + }), + ); + }); + it('should initialize with Unauthenticated state', async () => { - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); + // Because we defer refreshAuth, the initial state is safely caught here expect(result.current.authState).toBe(AuthState.Unauthenticated); - await waitFor(() => { - expect(result.current.authState).toBe(AuthState.Authenticated); + await act(async () => { + deferredRefreshAuth.resolve(); }); + + expect(result.current.authState).toBe(AuthState.Authenticated); }); it('should set error if no auth type is selected and no env key', async () => { - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(undefined), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toBe( - 'No authentication method selected.', - ); - expect(result.current.authState).toBe(AuthState.Updating); - }); + // This happens synchronously, no deferred promise + expect(result.current.authError).toBe( + 'No authentication method selected.', + ); + expect(result.current.authState).toBe(AuthState.Updating); }); it('should set error if no auth type is selected but env key exists', async () => { process.env['GEMINI_API_KEY'] = 'env-key'; - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(undefined), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toContain( - 'Existing API key detected (GEMINI_API_KEY)', - ); - expect(result.current.authState).toBe(AuthState.Updating); - }); + expect(result.current.authError).toContain( + 'Existing API key detected (GEMINI_API_KEY)', + ); + expect(result.current.authState).toBe(AuthState.Updating); }); it('should transition to AwaitingApiKeyInput if USE_GEMINI and no key found', async () => { - mockLoadApiKey.mockResolvedValue(null); - const { result } = renderHook(() => + let deferredLoadKey: { resolve: (k: string | null) => void }; + mockLoadApiKey.mockImplementation( + () => + new Promise((resolve) => { + deferredLoadKey = { resolve }; + }), + ); + + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.USE_GEMINI), mockConfig), ); - await waitFor(() => { - expect(result.current.authState).toBe(AuthState.AwaitingApiKeyInput); + await act(async () => { + deferredLoadKey.resolve(null); }); + + expect(result.current.authState).toBe(AuthState.AwaitingApiKeyInput); }); it('should authenticate if USE_GEMINI and key is found', async () => { - mockLoadApiKey.mockResolvedValue('stored-key'); - const { result } = renderHook(() => + let deferredLoadKey: { resolve: (k: string | null) => void }; + mockLoadApiKey.mockImplementation( + () => + new Promise((resolve) => { + deferredLoadKey = { resolve }; + }), + ); + + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.USE_GEMINI), mockConfig), ); - await waitFor(() => { - expect(mockConfig.refreshAuth).toHaveBeenCalledWith( - AuthType.USE_GEMINI, - ); - expect(result.current.authState).toBe(AuthState.Authenticated); - expect(result.current.apiKeyDefaultValue).toBe('stored-key'); + await act(async () => { + deferredLoadKey.resolve('stored-key'); }); + + await act(async () => { + deferredRefreshAuth.resolve(); + }); + + expect(mockConfig.refreshAuth).toHaveBeenCalledWith(AuthType.USE_GEMINI); + expect(result.current.authState).toBe(AuthState.Authenticated); + expect(result.current.apiKeyDefaultValue).toBe('stored-key'); }); it('should authenticate if USE_GEMINI and env key is found', async () => { - mockLoadApiKey.mockResolvedValue(null); process.env['GEMINI_API_KEY'] = 'env-key'; - const { result } = renderHook(() => + + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.USE_GEMINI), mockConfig), ); - await waitFor(() => { - expect(mockConfig.refreshAuth).toHaveBeenCalledWith( - AuthType.USE_GEMINI, - ); - expect(result.current.authState).toBe(AuthState.Authenticated); - expect(result.current.apiKeyDefaultValue).toBe('env-key'); + await act(async () => { + deferredRefreshAuth.resolve(); }); + + expect(mockConfig.refreshAuth).toHaveBeenCalledWith(AuthType.USE_GEMINI); + expect(result.current.authState).toBe(AuthState.Authenticated); + expect(result.current.apiKeyDefaultValue).toBe('env-key'); }); it('should prioritize env key over stored key when both are present', async () => { - mockLoadApiKey.mockResolvedValue('stored-key'); process.env['GEMINI_API_KEY'] = 'env-key'; - const { result } = renderHook(() => + + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.USE_GEMINI), mockConfig), ); - await waitFor(() => { - expect(mockConfig.refreshAuth).toHaveBeenCalledWith( - AuthType.USE_GEMINI, - ); - expect(result.current.authState).toBe(AuthState.Authenticated); - // The environment key should take precedence - expect(result.current.apiKeyDefaultValue).toBe('env-key'); + await act(async () => { + deferredRefreshAuth.resolve(); }); + + expect(mockConfig.refreshAuth).toHaveBeenCalledWith(AuthType.USE_GEMINI); + expect(result.current.authState).toBe(AuthState.Authenticated); + expect(result.current.apiKeyDefaultValue).toBe('env-key'); }); it('should set error if validation fails', async () => { mockValidateAuthMethod.mockReturnValue('Validation Failed'); - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toBe('Validation Failed'); - expect(result.current.authState).toBe(AuthState.Updating); - }); + expect(result.current.authError).toBe('Validation Failed'); + expect(result.current.authState).toBe(AuthState.Updating); }); it('should set error if GEMINI_DEFAULT_AUTH_TYPE is invalid', async () => { process.env['GEMINI_DEFAULT_AUTH_TYPE'] = 'INVALID_TYPE'; - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toContain( - 'Invalid value for GEMINI_DEFAULT_AUTH_TYPE', - ); - expect(result.current.authState).toBe(AuthState.Updating); - }); + expect(result.current.authError).toContain( + 'Invalid value for GEMINI_DEFAULT_AUTH_TYPE', + ); + expect(result.current.authState).toBe(AuthState.Updating); }); it('should authenticate successfully for valid auth type', async () => { - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); - await waitFor(() => { - expect(mockConfig.refreshAuth).toHaveBeenCalledWith( - AuthType.LOGIN_WITH_GOOGLE, - ); - expect(result.current.authState).toBe(AuthState.Authenticated); - expect(result.current.authError).toBeNull(); + await act(async () => { + deferredRefreshAuth.resolve(); }); + + expect(mockConfig.refreshAuth).toHaveBeenCalledWith( + AuthType.LOGIN_WITH_GOOGLE, + ); + expect(result.current.authState).toBe(AuthState.Authenticated); + expect(result.current.authError).toBeNull(); }); it('should handle refreshAuth failure', async () => { - (mockConfig.refreshAuth as Mock).mockRejectedValue( - new Error('Auth Failed'), - ); - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toContain('Failed to sign in'); - expect(result.current.authState).toBe(AuthState.Updating); + await act(async () => { + deferredRefreshAuth.reject(new Error('Auth Failed')); }); + + expect(result.current.authError).toContain('Failed to sign in'); + expect(result.current.authState).toBe(AuthState.Updating); }); it('should handle ProjectIdRequiredError without "Failed to login" prefix', async () => { const projectIdError = new ProjectIdRequiredError(); - (mockConfig.refreshAuth as Mock).mockRejectedValue(projectIdError); - const { result } = renderHook(() => + const { result } = await renderHook(() => useAuthCommand(createSettings(AuthType.LOGIN_WITH_GOOGLE), mockConfig), ); - await waitFor(() => { - expect(result.current.authError).toBe( - 'This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID env var. See https://goo.gle/gemini-cli-auth-docs#workspace-gca', - ); - expect(result.current.authError).not.toContain('Failed to login'); - expect(result.current.authState).toBe(AuthState.Updating); + await act(async () => { + deferredRefreshAuth.reject(projectIdError); }); + + expect(result.current.authError).toBe( + 'This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID env var. See https://goo.gle/gemini-cli-auth-docs#workspace-gca', + ); + expect(result.current.authError).not.toContain('Failed to login'); + expect(result.current.authState).toBe(AuthState.Updating); }); }); }); diff --git a/packages/cli/src/ui/components/AboutBox.test.tsx b/packages/cli/src/ui/components/AboutBox.test.tsx index 1db36b1f60..9115ca31c1 100644 --- a/packages/cli/src/ui/components/AboutBox.test.tsx +++ b/packages/cli/src/ui/components/AboutBox.test.tsx @@ -25,10 +25,9 @@ describe('AboutBox', () => { }; it('renders with required props', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('About Gemini CLI'); expect(output).toContain('1.0.0'); @@ -46,10 +45,9 @@ describe('AboutBox', () => { ['tier', 'Enterprise', 'Tier'], ])('renders optional prop %s', async (prop, value, label) => { const props = { ...defaultProps, [prop]: value }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain(label); expect(output).toContain(value); @@ -58,10 +56,9 @@ describe('AboutBox', () => { it('renders Auth Method with email when userEmail is provided', async () => { const props = { ...defaultProps, userEmail: 'test@example.com' }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('Signed in with Google (test@example.com)'); unmount(); @@ -69,10 +66,9 @@ describe('AboutBox', () => { it('renders Auth Method correctly when not oauth', async () => { const props = { ...defaultProps, selectedAuthType: 'api-key' }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('api-key'); unmount(); diff --git a/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx b/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx index 19db058b87..76a36fe4dc 100644 --- a/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx +++ b/packages/cli/src/ui/components/AdminSettingsChangedDialog.test.tsx @@ -17,15 +17,14 @@ describe('AdminSettingsChangedDialog', () => { }); it('renders correctly', async () => { - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('restarts on "r" key press', async () => { - const { stdin, waitUntilReady } = await renderWithProviders( + const { stdin } = await renderWithProviders( , { uiActions: { @@ -33,7 +32,6 @@ describe('AdminSettingsChangedDialog', () => { }, }, ); - await waitUntilReady(); act(() => { stdin.write('r'); @@ -43,7 +41,7 @@ describe('AdminSettingsChangedDialog', () => { }); it.each(['r', 'R'])('restarts on "%s" key press', async (key) => { - const { stdin, waitUntilReady } = await renderWithProviders( + const { stdin } = await renderWithProviders( , { uiActions: { @@ -51,7 +49,6 @@ describe('AdminSettingsChangedDialog', () => { }, }, ); - await waitUntilReady(); act(() => { stdin.write(key); diff --git a/packages/cli/src/ui/components/AgentConfigDialog.test.tsx b/packages/cli/src/ui/components/AgentConfigDialog.test.tsx index a2bfe052bb..2c6ea454db 100644 --- a/packages/cli/src/ui/components/AgentConfigDialog.test.tsx +++ b/packages/cli/src/ui/components/AgentConfigDialog.test.tsx @@ -126,7 +126,6 @@ describe('AgentConfigDialog', () => { />, { settings, uiState: { mainAreaWidth: 100 } }, ); - await result.waitUntilReady(); return result; }; diff --git a/packages/cli/src/ui/components/AlternateBufferQuittingDisplay.test.tsx b/packages/cli/src/ui/components/AlternateBufferQuittingDisplay.test.tsx index da71895485..571e0d36d3 100644 --- a/packages/cli/src/ui/components/AlternateBufferQuittingDisplay.test.tsx +++ b/packages/cli/src/ui/components/AlternateBufferQuittingDisplay.test.tsx @@ -108,7 +108,7 @@ describe('AlternateBufferQuittingDisplay', () => { it('renders with active and pending tool messages', async () => { persistentStateMock.setData({ tipsShown: 0 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -118,14 +118,13 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot('with_history_and_pending'); unmount(); }); it('renders with empty history and no pending items', async () => { persistentStateMock.setData({ tipsShown: 0 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -135,14 +134,13 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot('empty'); unmount(); }); it('renders with history but no pending items', async () => { persistentStateMock.setData({ tipsShown: 0 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -152,14 +150,13 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot('with_history_no_pending'); unmount(); }); it('renders with pending items but no history', async () => { persistentStateMock.setData({ tipsShown: 0 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -169,7 +166,6 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot('with_pending_no_history'); unmount(); }); @@ -195,7 +191,7 @@ describe('AlternateBufferQuittingDisplay', () => { ], }, ]; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -205,7 +201,6 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('Action Required (was prompted):'); expect(output).toContain('confirming_tool'); @@ -220,7 +215,7 @@ describe('AlternateBufferQuittingDisplay', () => { { id: 1, type: 'user', text: 'Hello Gemini' }, { id: 2, type: 'gemini', text: 'Hello User!' }, ]; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -230,7 +225,6 @@ describe('AlternateBufferQuittingDisplay', () => { }, }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot('with_user_gemini_messages'); unmount(); }); diff --git a/packages/cli/src/ui/components/AnsiOutput.test.tsx b/packages/cli/src/ui/components/AnsiOutput.test.tsx index ac824fefe6..758361be0a 100644 --- a/packages/cli/src/ui/components/AnsiOutput.test.tsx +++ b/packages/cli/src/ui/components/AnsiOutput.test.tsx @@ -29,10 +29,9 @@ describe('', () => { createAnsiToken({ text: 'world!' }), ], ]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame().trim()).toBe('Hello, world!'); unmount(); }); @@ -47,10 +46,9 @@ describe('', () => { { style: { inverse: true }, text: 'Inverse' }, ])('correctly applies style $text', async ({ style, text }) => { const data: AnsiOutput = [[createAnsiToken({ text, ...style })]]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame().trim()).toBe(text); unmount(); }); @@ -61,10 +59,9 @@ describe('', () => { { color: { fg: '#00ff00', bg: '#ff00ff' }, text: 'Green FG Magenta BG' }, ])('correctly applies color $text', async ({ color, text }) => { const data: AnsiOutput = [[createAnsiToken({ text, ...color })]]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame().trim()).toBe(text); unmount(); }); @@ -76,10 +73,9 @@ describe('', () => { [createAnsiToken({ text: 'Third line' })], [createAnsiToken({ text: '' })], ]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toBeDefined(); const lines = output.split('\n'); @@ -96,10 +92,9 @@ describe('', () => { [createAnsiToken({ text: 'Line 3' })], [createAnsiToken({ text: 'Line 4' })], ]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).not.toContain('Line 1'); expect(output).not.toContain('Line 2'); @@ -115,10 +110,9 @@ describe('', () => { [createAnsiToken({ text: 'Line 3' })], [createAnsiToken({ text: 'Line 4' })], ]; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).not.toContain('Line 1'); expect(output).not.toContain('Line 2'); @@ -135,7 +129,7 @@ describe('', () => { [createAnsiToken({ text: 'Line 4' })], ]; // availableTerminalHeight=3, maxLines=2 => show 2 lines - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { width={80} />, ); - await waitUntilReady(); const output = lastFrame(); expect(output).not.toContain('Line 2'); expect(output).toContain('Line 3'); @@ -156,10 +149,9 @@ describe('', () => { for (let i = 0; i < 1000; i++) { largeData.push([createAnsiToken({ text: `Line ${i}` })]); } - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); // We are just checking that it renders something without crashing. expect(lastFrame()).toBeDefined(); unmount(); diff --git a/packages/cli/src/ui/components/AppHeader.test.tsx b/packages/cli/src/ui/components/AppHeader.test.tsx index 0d7e2b3a7b..8ff4caaacf 100644 --- a/packages/cli/src/ui/components/AppHeader.test.tsx +++ b/packages/cli/src/ui/components/AppHeader.test.tsx @@ -27,13 +27,12 @@ describe('', () => { bannerVisible: true, }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).toContain('This is the default banner'); expect(lastFrame()).toMatchSnapshot(); @@ -50,13 +49,12 @@ describe('', () => { bannerVisible: true, }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).toContain('There are capacity issues'); expect(lastFrame()).toMatchSnapshot(); @@ -72,13 +70,12 @@ describe('', () => { }, }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).not.toContain('Banner'); expect(lastFrame()).toMatchSnapshot(); @@ -103,13 +100,12 @@ describe('', () => { }, }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).not.toContain('This is the default banner'); expect(lastFrame()).toMatchSnapshot(); @@ -129,13 +125,12 @@ describe('', () => { // and interfering with the expected persistentState.set call. persistentStateMock.setData({ tipsShown: 10 }); - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(persistentStateMock.set).toHaveBeenCalledWith( 'defaultBannerShownCount', @@ -159,13 +154,12 @@ describe('', () => { bannerVisible: true, }; - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).not.toContain('First line\\nSecond line'); unmount(); @@ -183,13 +177,12 @@ describe('', () => { persistentStateMock.setData({ tipsShown: 5 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).toContain('Tips'); expect(persistentStateMock.set).toHaveBeenCalledWith('tipsShown', 6); @@ -206,13 +199,12 @@ describe('', () => { persistentStateMock.setData({ tipsShown: 10 }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState, }, ); - await waitUntilReady(); expect(lastFrame()).not.toContain('Tips'); unmount(); @@ -234,7 +226,6 @@ describe('', () => { const session1 = await renderWithProviders(, { uiState, }); - await session1.waitUntilReady(); expect(session1.lastFrame()).toContain('Tips'); expect(persistentStateMock.get('tipsShown')).toBe(10); @@ -245,7 +236,6 @@ describe('', () => { , {}, ); - await session2.waitUntilReady(); expect(session2.lastFrame()).not.toContain('Tips'); session2.unmount(); diff --git a/packages/cli/src/ui/components/ApprovalModeIndicator.test.tsx b/packages/cli/src/ui/components/ApprovalModeIndicator.test.tsx index 4386891c7a..1b2decbe16 100644 --- a/packages/cli/src/ui/components/ApprovalModeIndicator.test.tsx +++ b/packages/cli/src/ui/components/ApprovalModeIndicator.test.tsx @@ -11,56 +11,50 @@ import { ApprovalMode } from '@google/gemini-cli-core'; describe('ApprovalModeIndicator', () => { it('renders correctly for AUTO_EDIT mode', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders correctly for AUTO_EDIT mode with plan enabled', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders correctly for PLAN mode', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders correctly for YOLO mode', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders correctly for DEFAULT mode', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders correctly for DEFAULT mode with plan enabled', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); }); diff --git a/packages/cli/src/ui/components/AskUserDialog.test.tsx b/packages/cli/src/ui/components/AskUserDialog.test.tsx index 8ed240389c..864800a061 100644 --- a/packages/cli/src/ui/components/AskUserDialog.test.tsx +++ b/packages/cli/src/ui/components/AskUserDialog.test.tsx @@ -48,7 +48,7 @@ describe('AskUserDialog', () => { ]; it('renders question and options', async () => { - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -397,7 +396,7 @@ describe('AskUserDialog', () => { }, ]; - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('hides progress header for single question', async () => { - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('shows keyboard hints', async () => { - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -471,7 +467,6 @@ describe('AskUserDialog', () => { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toContain('Which testing framework?'); writeKey(stdin, '\x1b[C'); // Right arrow @@ -582,7 +577,7 @@ describe('AskUserDialog', () => { }, ]; - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -736,7 +730,7 @@ describe('AskUserDialog', () => { }, ]; - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -759,7 +752,7 @@ describe('AskUserDialog', () => { }, ]; - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -820,7 +812,7 @@ describe('AskUserDialog', () => { }, ]; - const { lastFrame, waitUntilReady } = await renderWithProviders( + const { lastFrame } = await renderWithProviders( { { width: 120 }, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); diff --git a/packages/cli/src/ui/components/BackgroundShellDisplay.test.tsx b/packages/cli/src/ui/components/BackgroundShellDisplay.test.tsx index 847dcd9a87..c097028a0d 100644 --- a/packages/cli/src/ui/components/BackgroundShellDisplay.test.tsx +++ b/packages/cli/src/ui/components/BackgroundShellDisplay.test.tsx @@ -145,7 +145,7 @@ describe('', () => { it('renders the output of the active shell', async () => { const width = 80; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); @@ -166,7 +165,7 @@ describe('', () => { it('renders tabs for multiple shells', async () => { const width = 100; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); @@ -187,7 +185,7 @@ describe('', () => { it('highlights the focused state', async () => { const width = 80; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); @@ -208,7 +205,7 @@ describe('', () => { it('resizes the PTY on mount and when dimensions change', async () => { const width = 80; - const { rerender, waitUntilReady, unmount } = render( + const { rerender, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(ShellExecutionService.resizePty).toHaveBeenCalledWith( shell1.pid, @@ -241,7 +237,6 @@ describe('', () => { /> , ); - await waitUntilReady(); expect(ShellExecutionService.resizePty).toHaveBeenCalledWith( shell1.pid, @@ -253,7 +248,7 @@ describe('', () => { it('renders the process list when isListOpenProp is true', async () => { const width = 80; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); @@ -274,7 +268,7 @@ describe('', () => { it('selects the current process and closes the list when Ctrl+L is pressed in list view', async () => { const width = 80; - const { waitUntilReady, unmount } = render( + const { unmount } = await render( ', () => { , width, ); - await waitUntilReady(); // Simulate down arrow to select the second process (handled by RadioButtonSelect) await act(async () => { simulateKey({ name: 'down' }); }); - await waitUntilReady(); // Simulate Ctrl+L (handled by BackgroundShellDisplay) await act(async () => { simulateKey({ name: 'l', ctrl: true }); }); - await waitUntilReady(); expect(mockSetActiveBackgroundShellPid).toHaveBeenCalledWith(shell2.pid); expect(mockSetIsBackgroundShellListOpen).toHaveBeenCalledWith(false); @@ -308,7 +299,7 @@ describe('', () => { it('kills the highlighted process when Ctrl+K is pressed in list view', async () => { const width = 80; - const { waitUntilReady, unmount } = render( + const { unmount } = await render( ', () => { , width, ); - await waitUntilReady(); // Initial state: shell1 (active) is highlighted @@ -329,13 +319,11 @@ describe('', () => { await act(async () => { simulateKey({ name: 'down' }); }); - await waitUntilReady(); // Press Ctrl+K await act(async () => { simulateKey({ name: 'k', ctrl: true }); }); - await waitUntilReady(); expect(mockDismissBackgroundShell).toHaveBeenCalledWith(shell2.pid); unmount(); @@ -343,7 +331,7 @@ describe('', () => { it('kills the active process when Ctrl+K is pressed in output view', async () => { const width = 80; - const { waitUntilReady, unmount } = render( + const { unmount } = await render( ', () => { , width, ); - await waitUntilReady(); await act(async () => { simulateKey({ name: 'k', ctrl: true }); }); - await waitUntilReady(); expect(mockDismissBackgroundShell).toHaveBeenCalledWith(shell1.pid); unmount(); @@ -370,7 +356,7 @@ describe('', () => { it('scrolls to active shell when list opens', async () => { // shell2 is active const width = 80; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); @@ -402,7 +387,7 @@ describe('', () => { mockShells.set(exitedShell.pid, exitedShell); const width = 80; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( ', () => { , width, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); diff --git a/packages/cli/src/ui/components/Checklist.test.tsx b/packages/cli/src/ui/components/Checklist.test.tsx index 442ee0400f..329a560aec 100644 --- a/packages/cli/src/ui/components/Checklist.test.tsx +++ b/packages/cli/src/ui/components/Checklist.test.tsx @@ -18,10 +18,9 @@ describe('', () => { ]; it('renders nothing when list is empty', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame({ allowEmpty: true })).toBe(''); }); @@ -30,15 +29,14 @@ describe('', () => { { status: 'completed', label: 'Task 1' }, { status: 'cancelled', label: 'Task 2' }, ]; - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame({ allowEmpty: true })).toBe(''); }); it('renders summary view correctly (collapsed)', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( ', () => { toggleHint="toggle me" />, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('renders expanded view correctly', async () => { - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( ', () => { toggleHint="toggle me" />, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -68,10 +64,9 @@ describe('', () => { { status: 'completed', label: 'Task 1' }, { status: 'pending', label: 'Task 2' }, ]; - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); }); diff --git a/packages/cli/src/ui/components/ChecklistItem.test.tsx b/packages/cli/src/ui/components/ChecklistItem.test.tsx index 4176f7914b..c71af523e1 100644 --- a/packages/cli/src/ui/components/ChecklistItem.test.tsx +++ b/packages/cli/src/ui/components/ChecklistItem.test.tsx @@ -17,8 +17,7 @@ describe('', () => { { status: 'cancelled', label: 'Skipped this' }, { status: 'blocked', label: 'Blocked this' }, ] as ChecklistItemData[])('renders %s item correctly', async (item) => { - const { lastFrame, waitUntilReady } = render(); - await waitUntilReady(); + const { lastFrame } = await render(); expect(lastFrame()).toMatchSnapshot(); }); @@ -28,12 +27,11 @@ describe('', () => { label: 'This is a very long text that should be truncated because the wrap prop is set to truncate', }; - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); @@ -43,12 +41,11 @@ describe('', () => { label: 'This is a very long text that should wrap because the default behavior is wrapping', }; - const { lastFrame, waitUntilReady } = render( + const { lastFrame } = await render( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); }); diff --git a/packages/cli/src/ui/components/CliSpinner.test.tsx b/packages/cli/src/ui/components/CliSpinner.test.tsx index cca997f370..4da6abb199 100644 --- a/packages/cli/src/ui/components/CliSpinner.test.tsx +++ b/packages/cli/src/ui/components/CliSpinner.test.tsx @@ -17,10 +17,7 @@ describe('', () => { it('should increment debugNumAnimatedComponents on mount and decrement on unmount', async () => { expect(debugState.debugNumAnimatedComponents).toBe(0); - const { waitUntilReady, unmount } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { unmount } = await renderWithProviders(); expect(debugState.debugNumAnimatedComponents).toBe(1); unmount(); expect(debugState.debugNumAnimatedComponents).toBe(0); @@ -28,11 +25,9 @@ describe('', () => { it('should not render when showSpinner is false', async () => { const settings = createMockSettings({ ui: { showSpinner: false } }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( - , - { settings }, - ); - await waitUntilReady(); + const { lastFrame, unmount } = await renderWithProviders(, { + settings, + }); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); diff --git a/packages/cli/src/ui/components/ColorsDisplay.test.tsx b/packages/cli/src/ui/components/ColorsDisplay.test.tsx index fdd08fd653..d934831c0e 100644 --- a/packages/cli/src/ui/components/ColorsDisplay.test.tsx +++ b/packages/cli/src/ui/components/ColorsDisplay.test.tsx @@ -96,10 +96,9 @@ describe('ColorsDisplay', () => { it('renders correctly', async () => { const mockTheme = themeManager.getActiveTheme(); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); // Check for title and description diff --git a/packages/cli/src/ui/components/Composer.test.tsx b/packages/cli/src/ui/components/Composer.test.tsx index 641fc24810..8df5f690e7 100644 --- a/packages/cli/src/ui/components/Composer.test.tsx +++ b/packages/cli/src/ui/components/Composer.test.tsx @@ -251,7 +251,7 @@ const renderComposer = async ( config = createMockConfig(), uiActions = createMockUIActions(), ) => { - const result = render( + const result = await render( @@ -262,7 +262,6 @@ const renderComposer = async ( , ); - await result.waitUntilReady(); // Wait for shortcuts hint debounce if using fake timers if (vi.isFakeTimers()) { diff --git a/packages/cli/src/ui/components/ConfigInitDisplay.test.tsx b/packages/cli/src/ui/components/ConfigInitDisplay.test.tsx index 45ead4862e..b4ae8b93b1 100644 --- a/packages/cli/src/ui/components/ConfigInitDisplay.test.tsx +++ b/packages/cli/src/ui/components/ConfigInitDisplay.test.tsx @@ -43,10 +43,7 @@ describe('ConfigInitDisplay', () => { }); it('renders initial state', async () => { - const { lastFrame, waitUntilReady } = await renderWithProviders( - , - ); - await waitUntilReady(); + const { lastFrame } = await renderWithProviders(); expect(lastFrame()).toMatchSnapshot(); }); diff --git a/packages/cli/src/ui/components/ConsentPrompt.test.tsx b/packages/cli/src/ui/components/ConsentPrompt.test.tsx index dd69c44dd5..09a2dde16e 100644 --- a/packages/cli/src/ui/components/ConsentPrompt.test.tsx +++ b/packages/cli/src/ui/components/ConsentPrompt.test.tsx @@ -33,14 +33,13 @@ describe('ConsentPrompt', () => { it('renders a string prompt with MarkdownDisplay', async () => { const prompt = 'Are you sure?'; - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); expect(MockedMarkdownDisplay).toHaveBeenCalledWith( { @@ -55,14 +54,13 @@ describe('ConsentPrompt', () => { it('renders a ReactNode prompt directly', async () => { const prompt = Are you sure?; - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(MockedMarkdownDisplay).not.toHaveBeenCalled(); expect(lastFrame()).toContain('Are you sure?'); @@ -71,14 +69,13 @@ describe('ConsentPrompt', () => { it('calls onConfirm with true when "Yes" is selected', async () => { const prompt = 'Are you sure?'; - const { waitUntilReady, unmount } = render( + const { waitUntilReady, unmount } = await render( , ); - await waitUntilReady(); const onSelect = MockedRadioButtonSelect.mock.calls[0][0].onSelect; await act(async () => { @@ -92,14 +89,13 @@ describe('ConsentPrompt', () => { it('calls onConfirm with false when "No" is selected', async () => { const prompt = 'Are you sure?'; - const { waitUntilReady, unmount } = render( + const { waitUntilReady, unmount } = await render( , ); - await waitUntilReady(); const onSelect = MockedRadioButtonSelect.mock.calls[0][0].onSelect; await act(async () => { @@ -113,14 +109,13 @@ describe('ConsentPrompt', () => { it('passes correct items to RadioButtonSelect', async () => { const prompt = 'Are you sure?'; - const { waitUntilReady, unmount } = render( + const { unmount } = await render( , ); - await waitUntilReady(); expect(MockedRadioButtonSelect).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/cli/src/ui/components/ConsoleSummaryDisplay.test.tsx b/packages/cli/src/ui/components/ConsoleSummaryDisplay.test.tsx index cb8db1a895..b7662c3a26 100644 --- a/packages/cli/src/ui/components/ConsoleSummaryDisplay.test.tsx +++ b/packages/cli/src/ui/components/ConsoleSummaryDisplay.test.tsx @@ -10,10 +10,9 @@ import { describe, it, expect } from 'vitest'; describe('ConsoleSummaryDisplay', () => { it('renders nothing when errorCount is 0', async () => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -22,10 +21,9 @@ describe('ConsoleSummaryDisplay', () => { [1, '1 error'], [5, '5 errors'], ])('renders correct message for %i errors', async (count, expectedText) => { - const { lastFrame, waitUntilReady, unmount } = render( + const { lastFrame, unmount } = await render( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain(expectedText); expect(output).toContain('✖'); diff --git a/packages/cli/src/ui/components/ContextSummaryDisplay.test.tsx b/packages/cli/src/ui/components/ContextSummaryDisplay.test.tsx index f48cfb2a31..1049e97912 100644 --- a/packages/cli/src/ui/components/ContextSummaryDisplay.test.tsx +++ b/packages/cli/src/ui/components/ContextSummaryDisplay.test.tsx @@ -26,8 +26,7 @@ const renderWithWidth = async ( props: React.ComponentProps, ) => { useTerminalSizeMock.mockReturnValue({ columns: width, rows: 24 }); - const result = render(); - await result.waitUntilReady(); + const result = await render(); return result; }; diff --git a/packages/cli/src/ui/components/ContextUsageDisplay.test.tsx b/packages/cli/src/ui/components/ContextUsageDisplay.test.tsx index 904e06635c..d8ec1650ee 100644 --- a/packages/cli/src/ui/components/ContextUsageDisplay.test.tsx +++ b/packages/cli/src/ui/components/ContextUsageDisplay.test.tsx @@ -19,35 +19,33 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { describe('ContextUsageDisplay', () => { it('renders correct percentage used', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('50% used'); unmount(); }); it('renders correctly when usage is 0%', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('0% used'); unmount(); }); it('renders abbreviated label when terminal width is small', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { />, { width: 80 }, ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('20%'); expect(output).not.toContain('context used'); @@ -63,28 +60,26 @@ describe('ContextUsageDisplay', () => { }); it('renders 80% correctly', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('80% used'); unmount(); }); it('renders 100% when full', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('100% used'); unmount(); diff --git a/packages/cli/src/ui/components/CopyModeWarning.test.tsx b/packages/cli/src/ui/components/CopyModeWarning.test.tsx index 6f202ced4a..cc20a142dd 100644 --- a/packages/cli/src/ui/components/CopyModeWarning.test.tsx +++ b/packages/cli/src/ui/components/CopyModeWarning.test.tsx @@ -22,8 +22,7 @@ describe('CopyModeWarning', () => { mockUseUIState.mockReturnValue({ copyModeEnabled: false, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -32,8 +31,7 @@ describe('CopyModeWarning', () => { mockUseUIState.mockReturnValue({ copyModeEnabled: true, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame()).toContain('In Copy Mode'); expect(lastFrame()).toContain('Use Page Up/Down to scroll'); expect(lastFrame()).toContain('Press Ctrl+S or any other key to exit'); diff --git a/packages/cli/src/ui/components/DebugProfiler.test.tsx b/packages/cli/src/ui/components/DebugProfiler.test.tsx index d4c0e28902..a014c740f0 100644 --- a/packages/cli/src/ui/components/DebugProfiler.test.tsx +++ b/packages/cli/src/ui/components/DebugProfiler.test.tsx @@ -242,8 +242,7 @@ describe('DebugProfiler Component', () => { showDebugProfiler: false, constrainHeight: false, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -257,8 +256,7 @@ describe('DebugProfiler Component', () => { profiler.totalIdleFrames = 5; profiler.totalFlickerFrames = 2; - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); const output = lastFrame(); expect(output).toContain('Renders: 10 (total)'); @@ -275,8 +273,7 @@ describe('DebugProfiler Component', () => { const reportActionSpy = vi.spyOn(profiler, 'reportAction'); - const { waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { waitUntilReady, unmount } = await render(); await act(async () => { coreEvents.emitModelChanged('new-model'); @@ -295,8 +292,7 @@ describe('DebugProfiler Component', () => { const reportActionSpy = vi.spyOn(profiler, 'reportAction'); - const { waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { waitUntilReady, unmount } = await render(); await act(async () => { appEvents.emit(AppEvent.SelectionWarning); diff --git a/packages/cli/src/ui/components/DetailedMessagesDisplay.test.tsx b/packages/cli/src/ui/components/DetailedMessagesDisplay.test.tsx index b2f4185842..30f98a6eda 100644 --- a/packages/cli/src/ui/components/DetailedMessagesDisplay.test.tsx +++ b/packages/cli/src/ui/components/DetailedMessagesDisplay.test.tsx @@ -41,13 +41,12 @@ describe('DetailedMessagesDisplay', () => { }); }); it('renders nothing when messages are empty', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { settings: createMockSettings({ ui: { errorVerbosity: 'full' } }), }, ); - await waitUntilReady(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -64,13 +63,12 @@ describe('DetailedMessagesDisplay', () => { clearConsoleMessages: vi.fn(), }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { settings: createMockSettings({ ui: { errorVerbosity: 'full' } }), }, ); - await waitUntilReady(); const output = lastFrame(); expect(output).toMatchSnapshot(); @@ -86,13 +84,12 @@ describe('DetailedMessagesDisplay', () => { clearConsoleMessages: vi.fn(), }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { settings: createMockSettings({ ui: { errorVerbosity: 'low' } }), }, ); - await waitUntilReady(); expect(lastFrame()).toContain('(F12 to close)'); unmount(); }); @@ -106,13 +103,12 @@ describe('DetailedMessagesDisplay', () => { clearConsoleMessages: vi.fn(), }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { settings: createMockSettings({ ui: { errorVerbosity: 'full' } }), }, ); - await waitUntilReady(); expect(lastFrame()).toContain('(F12 to close)'); unmount(); }); @@ -126,13 +122,12 @@ describe('DetailedMessagesDisplay', () => { clearConsoleMessages: vi.fn(), }); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { settings: createMockSettings({ ui: { errorVerbosity: 'full' } }), }, ); - await waitUntilReady(); const output = lastFrame(); expect(output).toMatchSnapshot(); diff --git a/packages/cli/src/ui/components/DialogManager.test.tsx b/packages/cli/src/ui/components/DialogManager.test.tsx index 6f6dbb0289..31b28f5223 100644 --- a/packages/cli/src/ui/components/DialogManager.test.tsx +++ b/packages/cli/src/ui/components/DialogManager.test.tsx @@ -104,11 +104,10 @@ describe('DialogManager', () => { }; it('renders nothing by default', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: baseUiState as Partial as UIState }, ); - await waitUntilReady(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -197,7 +196,7 @@ describe('DialogManager', () => { it.each(testCases)( 'renders %s when state is %o', async (uiStateOverride, expectedComponent) => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , { uiState: { @@ -206,7 +205,6 @@ describe('DialogManager', () => { } as Partial as UIState, }, ); - await waitUntilReady(); expect(lastFrame()).toContain(expectedComponent); unmount(); }, diff --git a/packages/cli/src/ui/components/EditorSettingsDialog.test.tsx b/packages/cli/src/ui/components/EditorSettingsDialog.test.tsx index bd995652b1..18b47def7b 100644 --- a/packages/cli/src/ui/components/EditorSettingsDialog.test.tsx +++ b/packages/cli/src/ui/components/EditorSettingsDialog.test.tsx @@ -55,27 +55,25 @@ describe('EditorSettingsDialog', () => { renderWithProviders(ui); it('renders correctly', async () => { - const { lastFrame, waitUntilReady } = await renderWithProvider( + const { lastFrame } = await renderWithProvider( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); }); it('calls onSelect when an editor is selected', async () => { const onSelect = vi.fn(); - const { lastFrame, waitUntilReady } = await renderWithProvider( + const { lastFrame } = await renderWithProvider( , ); - await waitUntilReady(); expect(lastFrame()).toContain('VS Code'); }); @@ -88,7 +86,6 @@ describe('EditorSettingsDialog', () => { onExit={vi.fn()} />, ); - await waitUntilReady(); // Initial focus on editor expect(lastFrame()).toContain('> Select Editor'); @@ -134,7 +131,6 @@ describe('EditorSettingsDialog', () => { onExit={onExit} />, ); - await waitUntilReady(); await act(async () => { stdin.write('\u001B'); // Escape @@ -162,14 +158,13 @@ describe('EditorSettingsDialog', () => { }, } as unknown as LoadedSettings; - const { lastFrame, waitUntilReady } = await renderWithProvider( + const { lastFrame } = await renderWithProvider( , ); - await waitUntilReady(); const frame = lastFrame() || ''; if (!frame.includes('(Also modified')) { diff --git a/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx b/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx index 23a2038b10..74de1a8a41 100644 --- a/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx +++ b/packages/cli/src/ui/components/EmptyWalletDialog.test.tsx @@ -30,7 +30,7 @@ describe('EmptyWalletDialog', () => { describe('rendering', () => { it('should match snapshot with fallback available', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { onChoice={mockOnChoice} />, ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('should match snapshot without fallback', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toMatchSnapshot(); unmount(); }); it('should display the model name and usage limit message', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame() ?? ''; expect(output).toContain('gemini-2.5-pro'); @@ -73,13 +70,12 @@ describe('EmptyWalletDialog', () => { }); it('should display purchase prompt and credits update notice', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame() ?? ''; expect(output).toContain('purchase more AI Credits'); @@ -90,14 +86,13 @@ describe('EmptyWalletDialog', () => { }); it('should display reset time when provided', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame() ?? ''; expect(output).toContain('3:45 PM'); @@ -106,13 +101,12 @@ describe('EmptyWalletDialog', () => { }); it('should not display reset time when not provided', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame() ?? ''; expect(output).not.toContain('Access resets at'); @@ -120,13 +114,12 @@ describe('EmptyWalletDialog', () => { }); it('should display slash command hints', async () => { - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); const output = lastFrame() ?? ''; expect(output).toContain('/stats'); @@ -139,14 +132,13 @@ describe('EmptyWalletDialog', () => { describe('onChoice handling', () => { it('should call onGetCredits and onChoice when get_credits is selected', async () => { // get_credits is the first item, so just press Enter - const { unmount, stdin, waitUntilReady } = await renderWithProviders( + const { unmount, stdin } = await renderWithProviders( , ); - await waitUntilReady(); writeKey(stdin, '\r'); @@ -158,13 +150,12 @@ describe('EmptyWalletDialog', () => { }); it('should call onChoice without onGetCredits when onGetCredits is not provided', async () => { - const { unmount, stdin, waitUntilReady } = await renderWithProviders( + const { unmount, stdin } = await renderWithProviders( , ); - await waitUntilReady(); writeKey(stdin, '\r'); @@ -177,14 +168,13 @@ describe('EmptyWalletDialog', () => { it('should call onChoice with use_fallback when selected', async () => { // With fallback: items are [get_credits, use_fallback, stop] // use_fallback is the second item: Down + Enter - const { unmount, stdin, waitUntilReady } = await renderWithProviders( + const { unmount, stdin } = await renderWithProviders( , ); - await waitUntilReady(); writeKey(stdin, '\x1b[B'); // Down arrow writeKey(stdin, '\r'); @@ -198,13 +188,12 @@ describe('EmptyWalletDialog', () => { it('should call onChoice with stop when selected', async () => { // Without fallback: items are [get_credits, stop] // stop is the second item: Down + Enter - const { unmount, stdin, waitUntilReady } = await renderWithProviders( + const { unmount, stdin } = await renderWithProviders( , ); - await waitUntilReady(); writeKey(stdin, '\x1b[B'); // Down arrow writeKey(stdin, '\r'); diff --git a/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx b/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx index f369e7ff8e..d6fc23dd70 100644 --- a/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx +++ b/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx @@ -440,36 +440,38 @@ Implement a comprehensive authentication system with multiple providers. return <>{children}; }; - const { stdin, lastFrame } = await renderWithProviders( - - - , - { - config: { - getTargetDir: () => mockTargetDir, - getIdeMode: () => false, - isTrustedFolder: () => true, - storage: { - getPlansDir: () => mockPlansDir, - }, - getFileSystemService: (): FileSystemService => ({ - readTextFile: vi.fn(), - writeTextFile: vi.fn(), + const { stdin, lastFrame } = await act(async () => + renderWithProviders( + + + , + { + config: { + getTargetDir: () => mockTargetDir, + getIdeMode: () => false, + isTrustedFolder: () => true, + storage: { + getPlansDir: () => mockPlansDir, + }, + getFileSystemService: (): FileSystemService => ({ + readTextFile: vi.fn(), + writeTextFile: vi.fn(), + }), + getUseAlternateBuffer: () => useAlternateBuffer ?? true, + } as unknown as import('@google/gemini-cli-core').Config, + settings: createMockSettings({ + ui: { useAlternateBuffer: useAlternateBuffer ?? true }, }), - getUseAlternateBuffer: () => useAlternateBuffer ?? true, - } as unknown as import('@google/gemini-cli-core').Config, - settings: createMockSettings({ - ui: { useAlternateBuffer: useAlternateBuffer ?? true }, - }), - }, + }, + ), ); await act(async () => { diff --git a/packages/cli/src/ui/components/ExitWarning.test.tsx b/packages/cli/src/ui/components/ExitWarning.test.tsx index 6d495a5e21..a504670d03 100644 --- a/packages/cli/src/ui/components/ExitWarning.test.tsx +++ b/packages/cli/src/ui/components/ExitWarning.test.tsx @@ -24,8 +24,7 @@ describe('ExitWarning', () => { ctrlCPressedOnce: false, ctrlDPressedOnce: false, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); @@ -36,8 +35,7 @@ describe('ExitWarning', () => { ctrlCPressedOnce: true, ctrlDPressedOnce: false, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame()).toContain('Press Ctrl+C again to exit'); unmount(); }); @@ -48,8 +46,7 @@ describe('ExitWarning', () => { ctrlCPressedOnce: false, ctrlDPressedOnce: true, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame()).toContain('Press Ctrl+D again to exit'); unmount(); }); @@ -60,8 +57,7 @@ describe('ExitWarning', () => { ctrlCPressedOnce: true, ctrlDPressedOnce: true, } as unknown as UIState); - const { lastFrame, waitUntilReady, unmount } = render(); - await waitUntilReady(); + const { lastFrame, unmount } = await render(); expect(lastFrame({ allowEmpty: true })).toBe(''); unmount(); }); diff --git a/packages/cli/src/ui/components/FolderTrustDialog.test.tsx b/packages/cli/src/ui/components/FolderTrustDialog.test.tsx index c1d04b3ff9..de6e8096ec 100644 --- a/packages/cli/src/ui/components/FolderTrustDialog.test.tsx +++ b/packages/cli/src/ui/components/FolderTrustDialog.test.tsx @@ -48,10 +48,9 @@ describe('FolderTrustDialog', () => { }); it('should render the dialog with title and description', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Do you trust the files in this folder?'); expect(lastFrame()).toContain( @@ -72,7 +71,7 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { }, ); - await waitUntilReady(); expect(lastFrame()).toContain('This folder contains:'); expect(lastFrame()).toContain('hidden'); unmount(); @@ -103,7 +101,7 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { }, ); - await waitUntilReady(); // With maxHeight=4, the intro text (4 lines) will take most of the space. // The discovery results will likely be hidden. expect(lastFrame()).toContain('hidden'); @@ -135,7 +132,7 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { }, ); - await waitUntilReady(); expect(lastFrame()).toContain('hidden'); unmount(); }); @@ -182,9 +178,7 @@ describe('FolderTrustDialog', () => { // Initial state: truncated await waitFor(() => { expect(lastFrame()).toContain('Do you trust the files in this folder?'); - // In standard terminal mode, the expansion hint is handled globally by ToastDisplay - // via AppContainer, so it should not be present in the dialog's local frame. - expect(lastFrame()).not.toContain('Press Ctrl+O'); + expect(lastFrame()).toContain('Press Ctrl+O'); expect(lastFrame()).toContain('hidden'); }); @@ -221,7 +215,6 @@ describe('FolderTrustDialog', () => { await renderWithProviders( , ); - await waitUntilReady(); await act(async () => { stdin.write('\u001b[27u'); // Press kitty escape key @@ -246,10 +239,9 @@ describe('FolderTrustDialog', () => { }); it('should display restart message when isRestarting is true', async () => { - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Gemini CLI is restarting'); unmount(); @@ -260,10 +252,9 @@ describe('FolderTrustDialog', () => { const relaunchApp = vi .spyOn(processUtils, 'relaunchApp') .mockResolvedValue(undefined); - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); await vi.advanceTimersByTimeAsync(250); expect(relaunchApp).toHaveBeenCalled(); unmount(); @@ -275,10 +266,9 @@ describe('FolderTrustDialog', () => { const relaunchApp = vi .spyOn(processUtils, 'relaunchApp') .mockResolvedValue(undefined); - const { waitUntilReady, unmount } = await renderWithProviders( + const { unmount } = await renderWithProviders( , ); - await waitUntilReady(); // Unmount immediately (before 250ms) unmount(); @@ -292,7 +282,6 @@ describe('FolderTrustDialog', () => { const { stdin, waitUntilReady, unmount } = await renderWithProviders( , ); - await waitUntilReady(); await act(async () => { stdin.write('r'); @@ -308,30 +297,27 @@ describe('FolderTrustDialog', () => { describe('directory display', () => { it('should correctly display the folder name for a nested directory', async () => { mockedCwd.mockReturnValue('/home/user/project'); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Trust folder (project)'); unmount(); }); it('should correctly display the parent folder name for a nested directory', async () => { mockedCwd.mockReturnValue('/home/user/project'); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Trust parent folder (user)'); unmount(); }); it('should correctly display an empty parent folder name for a directory directly under root', async () => { mockedCwd.mockReturnValue('/project'); - const { lastFrame, waitUntilReady, unmount } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Trust parent folder ()'); unmount(); }); @@ -348,7 +334,7 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { { width: 80 }, ); - await waitUntilReady(); expect(lastFrame()).toContain('This folder contains:'); expect(lastFrame()).toContain('• Commands (2):'); expect(lastFrame()).toContain('- cmd1'); @@ -386,14 +371,13 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: ['Dangerous setting detected!'], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Security Warnings:'); expect(lastFrame()).toContain('Dangerous setting detected!'); unmount(); @@ -410,14 +394,13 @@ describe('FolderTrustDialog', () => { discoveryErrors: ['Failed to load custom commands'], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( , ); - await waitUntilReady(); expect(lastFrame()).toContain('Discovery Errors:'); expect(lastFrame()).toContain('Failed to load custom commands'); unmount(); @@ -434,7 +417,7 @@ describe('FolderTrustDialog', () => { discoveryErrors: [], securityWarnings: [], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { }, ); - await waitUntilReady(); // In alternate buffer + expanded, the title should be visible (StickyHeader) expect(lastFrame()).toContain('Do you trust the files in this folder?'); // And it should NOT use MaxSizedBox truncation @@ -470,7 +452,7 @@ describe('FolderTrustDialog', () => { securityWarnings: [`${ansiRed}warning-with-ansi${ansiReset}`], }; - const { lastFrame, unmount, waitUntilReady } = await renderWithProviders( + const { lastFrame, unmount } = await renderWithProviders( { { width: 100, uiState: { terminalHeight: 40 } }, ); - await waitUntilReady(); const output = lastFrame(); expect(output).toContain('cmd-with-ansi'); diff --git a/packages/cli/src/ui/components/Footer.test.tsx b/packages/cli/src/ui/components/Footer.test.tsx index 39f20e1c86..c0a52af868 100644 --- a/packages/cli/src/ui/components/Footer.test.tsx +++ b/packages/cli/src/ui/components/Footer.test.tsx @@ -138,33 +138,25 @@ describe('