Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions packages/core/src/code_assist/oauth2.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import { writeToStdout } from '../utils/stdio.js';
import { FatalCancellationError } from '../utils/errors.js';
import process from 'node:process';
import { coreEvents } from '../utils/events.js';
import { isHeadlessMode } from '../utils/headless.js';

vi.mock('node:os', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:os')>();
Expand All @@ -54,6 +55,9 @@ vi.mock('http');
vi.mock('open');
vi.mock('crypto');
vi.mock('node:readline');
vi.mock('../utils/headless.js', () => ({
isHeadlessMode: vi.fn(),
}));
vi.mock('../utils/browser.js', () => ({
shouldAttemptBrowserLaunch: () => true,
}));
Expand Down Expand Up @@ -98,6 +102,12 @@ global.fetch = vi.fn();

describe('oauth2', () => {
beforeEach(() => {
vi.mocked(isHeadlessMode).mockReturnValue(false);
(readline.createInterface as Mock).mockReturnValue({
question: vi.fn((_query, callback) => callback('')),
close: vi.fn(),
on: vi.fn(),
});
vi.spyOn(coreEvents, 'listenerCount').mockReturnValue(1);
vi.spyOn(coreEvents, 'emitConsentRequest').mockImplementation((payload) => {
payload.onConfirm(true);
Expand Down
17 changes: 17 additions & 0 deletions packages/core/src/mcp/oauth-provider.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,23 @@ vi.mock('../utils/events.js', () => ({
vi.mock('../utils/authConsent.js', () => ({
getConsentForOauth: vi.fn(() => Promise.resolve(true)),
}));
vi.mock('../utils/headless.js', () => ({
isHeadlessMode: vi.fn(() => false),
}));
vi.mock('node:readline', () => ({
default: {
createInterface: vi.fn(() => ({
question: vi.fn((_query, callback) => callback('')),
close: vi.fn(),
on: vi.fn(),
})),
},
createInterface: vi.fn(() => ({
question: vi.fn((_query, callback) => callback('')),
close: vi.fn(),
on: vi.fn(),
})),
}));

import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import * as http from 'node:http';
Expand Down
159 changes: 101 additions & 58 deletions packages/core/src/utils/authConsent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { describe, it, expect, vi } from 'vitest';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import type { Mock } from 'vitest';
import readline from 'node:readline';
import process from 'node:process';
Expand All @@ -27,73 +27,116 @@ vi.mock('./stdio.js', () => ({
}));

describe('getConsentForOauth', () => {
it('should use coreEvents when listeners are present', async () => {
beforeEach(() => {
vi.restoreAllMocks();
const mockEmitConsentRequest = vi.spyOn(coreEvents, 'emitConsentRequest');
const mockListenerCount = vi
.spyOn(coreEvents, 'listenerCount')
.mockReturnValue(1);
});

mockEmitConsentRequest.mockImplementation((payload) => {
payload.onConfirm(true);
describe('in interactive mode', () => {
beforeEach(() => {
(isHeadlessMode as Mock).mockReturnValue(false);
});

const result = await getConsentForOauth('Login required.');
it('should emit consent request when UI listeners are present', async () => {
const mockEmitConsentRequest = vi.spyOn(coreEvents, 'emitConsentRequest');
vi.spyOn(coreEvents, 'listenerCount').mockReturnValue(1);

expect(result).toBe(true);
expect(mockEmitConsentRequest).toHaveBeenCalledWith(
expect.objectContaining({
prompt: expect.stringContaining(
'Login required. Opening authentication page in your browser.',
),
}),
);
mockEmitConsentRequest.mockImplementation((payload) => {
payload.onConfirm(true);
});

mockListenerCount.mockRestore();
mockEmitConsentRequest.mockRestore();
});
const result = await getConsentForOauth('Login required.');

it('should use readline when no listeners are present and not headless', async () => {
vi.restoreAllMocks();
const mockListenerCount = vi
.spyOn(coreEvents, 'listenerCount')
.mockReturnValue(0);
(isHeadlessMode as Mock).mockReturnValue(false);

const mockReadline = {
on: vi.fn((event, callback) => {
if (event === 'line') {
callback('y');
}
}),
close: vi.fn(),
};
(readline.createInterface as Mock).mockReturnValue(mockReadline);

const result = await getConsentForOauth('Login required.');

expect(result).toBe(true);
expect(readline.createInterface).toHaveBeenCalled();
expect(writeToStdout).toHaveBeenCalledWith(
expect.stringContaining(
'Login required. Opening authentication page in your browser.',
),
);

mockListenerCount.mockRestore();
expect(result).toBe(true);
expect(mockEmitConsentRequest).toHaveBeenCalledWith(
expect.objectContaining({
prompt: expect.stringContaining(
'Login required. Opening authentication page in your browser.',
),
}),
);
});

it('should return false when user declines via UI', async () => {
const mockEmitConsentRequest = vi.spyOn(coreEvents, 'emitConsentRequest');
vi.spyOn(coreEvents, 'listenerCount').mockReturnValue(1);

mockEmitConsentRequest.mockImplementation((payload) => {
payload.onConfirm(false);
});

const result = await getConsentForOauth('Login required.');

expect(result).toBe(false);
});

it('should throw FatalAuthenticationError when no UI listeners are present', async () => {
vi.spyOn(coreEvents, 'listenerCount').mockReturnValue(0);

await expect(getConsentForOauth('Login required.')).rejects.toThrow(
FatalAuthenticationError,
);
});
});

it('should throw FatalAuthenticationError when no listeners and headless', async () => {
vi.restoreAllMocks();
const mockListenerCount = vi
.spyOn(coreEvents, 'listenerCount')
.mockReturnValue(0);
(isHeadlessMode as Mock).mockReturnValue(true);
describe('in non-interactive mode', () => {
beforeEach(() => {
(isHeadlessMode as Mock).mockReturnValue(true);
});

it('should use readline to prompt for consent', async () => {
const mockReadline = {
on: vi.fn((event, callback) => {
if (event === 'line') {
callback('y');
}
}),
close: vi.fn(),
};
(readline.createInterface as Mock).mockReturnValue(mockReadline);

const result = await getConsentForOauth('Login required.');

expect(result).toBe(true);
expect(readline.createInterface).toHaveBeenCalledWith(
expect.objectContaining({
terminal: true,
}),
);
expect(writeToStdout).toHaveBeenCalledWith(
expect.stringContaining('Login required.'),
);
});

await expect(getConsentForOauth('Login required.')).rejects.toThrow(
FatalAuthenticationError,
);
it('should accept empty response as "yes"', async () => {
const mockReadline = {
on: vi.fn((event, callback) => {
if (event === 'line') {
callback('');
}
}),
close: vi.fn(),
};
(readline.createInterface as Mock).mockReturnValue(mockReadline);

mockListenerCount.mockRestore();
const result = await getConsentForOauth('Login required.');

expect(result).toBe(true);
});

it('should return false when user declines via readline', async () => {
const mockReadline = {
on: vi.fn((event, callback) => {
if (event === 'line') {
callback('n');
}
}),
close: vi.fn(),
};
(readline.createInterface as Mock).mockReturnValue(mockReadline);

const result = await getConsentForOauth('Login required.');

expect(result).toBe(false);
});
});
});
19 changes: 9 additions & 10 deletions packages/core/src/utils/authConsent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,22 +12,21 @@ import { isHeadlessMode } from './headless.js';

/**
* Requests consent from the user for OAuth login.
* Handles both TTY and non-TTY environments.
* Handles both interactive and non-interactive (headless) modes.
*/
export async function getConsentForOauth(prompt: string): Promise<boolean> {
const finalPrompt = prompt + ' Opening authentication page in your browser. ';

if (coreEvents.listenerCount(CoreEvent.ConsentRequest) === 0) {
if (isHeadlessMode()) {
throw new FatalAuthenticationError(
'Interactive consent could not be obtained.\n' +
'Please run Gemini CLI in an interactive terminal to authenticate, or use NO_BROWSER=true for manual authentication.',
);
}
if (isHeadlessMode()) {
return getOauthConsentNonInteractive(finalPrompt);
} else if (coreEvents.listenerCount(CoreEvent.ConsentRequest) > 0) {
return getOauthConsentInteractive(finalPrompt);
}

return getOauthConsentInteractive(finalPrompt);
throw new FatalAuthenticationError(
'Authentication consent could not be obtained.\n' +
'Please run Gemini CLI in an interactive terminal to authenticate, ' +
'or use NO_BROWSER=true for manual authentication.',
);
}

async function getOauthConsentNonInteractive(prompt: string) {
Expand Down
Loading