feat: first time automated ollama install experience and openrouter (#3881)

Co-authored-by: Zane Staggs <zane@squareup.com>
This commit is contained in:
Michael Neale
2025-08-08 11:41:53 +10:00
committed by GitHub
parent 43d6d1a492
commit e124bb084f
9 changed files with 1444 additions and 10 deletions
@@ -0,0 +1,277 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { render, screen, waitFor, fireEvent } from '@testing-library/react';
import { OllamaSetup } from './OllamaSetup';
import * as ollamaDetection from '../utils/ollamaDetection';
import * as providerUtils from '../utils/providerUtils';
import { toastService } from '../toasts';
// Mock dependencies
vi.mock('../utils/ollamaDetection');
vi.mock('../utils/providerUtils');
vi.mock('../toasts');
// Mock useConfig hook
const mockUpsert = vi.fn();
const mockAddExtension = vi.fn();
const mockGetExtensions = vi.fn();
vi.mock('./ConfigContext', () => ({
useConfig: () => ({
upsert: mockUpsert,
addExtension: mockAddExtension,
getExtensions: mockGetExtensions,
}),
}));
describe('OllamaSetup', () => {
const mockOnSuccess = vi.fn();
const mockOnCancel = vi.fn();
beforeEach(() => {
vi.clearAllMocks();
// Default mocks
vi.mocked(ollamaDetection.getPreferredModel).mockReturnValue('gpt-oss:20b');
vi.mocked(ollamaDetection.getOllamaDownloadUrl).mockReturnValue('https://ollama.com/download');
});
describe('when Ollama is not detected', () => {
beforeEach(() => {
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: false,
host: 'http://127.0.0.1:11434',
});
});
it('should show installation instructions', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
expect(screen.getByText('Ollama Setup')).toBeInTheDocument();
expect(screen.getByText(/Ollama is not detected on your system/)).toBeInTheDocument();
});
});
it('should provide download link', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
const downloadLink = screen.getByRole('link', { name: /Install Ollama/ });
expect(downloadLink).toHaveAttribute('href', 'https://ollama.com/download');
expect(downloadLink).toHaveAttribute('target', '_blank');
});
});
it('should show polling state when install link is clicked', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
// Mock pollForOllama
const mockStopPolling = vi.fn();
vi.mocked(ollamaDetection.pollForOllama).mockReturnValue(mockStopPolling);
await waitFor(() => {
const installLink = screen.getByText('Install Ollama');
fireEvent.click(installLink);
});
expect(screen.getByText(/Waiting for Ollama to start/)).toBeInTheDocument();
expect(ollamaDetection.pollForOllama).toHaveBeenCalled();
});
it('should handle cancel button', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText('Use a different provider'));
});
expect(mockOnCancel).toHaveBeenCalled();
});
});
describe('when Ollama is detected but model is not available', () => {
beforeEach(() => {
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: true,
host: 'http://127.0.0.1:11434',
});
vi.mocked(ollamaDetection.hasModel).mockResolvedValue(false);
});
it('should show model download prompt', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
expect(screen.getByText(/The gpt-oss:20b model is not installed/)).toBeInTheDocument();
expect(screen.getByText(/Download gpt-oss:20b/)).toBeInTheDocument();
});
});
it('should handle model download', async () => {
vi.mocked(ollamaDetection.pullOllamaModel).mockResolvedValue(true);
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText(/Download gpt-oss:20b/));
});
await waitFor(() => {
expect(toastService.success).toHaveBeenCalledWith(
expect.objectContaining({
title: 'Model Downloaded!',
})
);
});
});
it('should handle download failure', async () => {
vi.mocked(ollamaDetection.pullOllamaModel).mockResolvedValue(false);
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText(/Download gpt-oss:20b/));
});
await waitFor(() => {
expect(toastService.error).toHaveBeenCalledWith(
expect.objectContaining({
title: 'Download Failed',
})
);
});
});
});
describe('when Ollama and model are both available', () => {
beforeEach(() => {
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: true,
host: 'http://127.0.0.1:11434',
});
vi.mocked(ollamaDetection.hasModel).mockResolvedValue(true);
});
it('should show ready state and connect button', async () => {
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
expect(screen.getByText(/Ollama is running on your system/)).toBeInTheDocument();
});
});
it('should handle successful connection', async () => {
vi.mocked(providerUtils.initializeSystem).mockResolvedValue(undefined);
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText(/Use Goose with Ollama/));
});
await waitFor(() => {
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_PROVIDER', 'ollama', false);
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_MODEL', 'gpt-oss:20b', false);
expect(mockUpsert).toHaveBeenCalledWith('OLLAMA_HOST', 'localhost', false);
expect(providerUtils.initializeSystem).toHaveBeenCalledWith(
'ollama',
'gpt-oss:20b',
expect.any(Object)
);
expect(toastService.success).toHaveBeenCalled();
expect(mockOnSuccess).toHaveBeenCalled();
});
});
it('should handle connection failure', async () => {
const testError = new Error('Initialization failed');
vi.mocked(providerUtils.initializeSystem).mockRejectedValue(testError);
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText('Use Goose with Ollama'));
});
await waitFor(() => {
expect(toastService.error).toHaveBeenCalledWith(
expect.objectContaining({
title: 'Connection Failed',
msg: expect.stringContaining('Initialization failed'),
})
);
});
});
});
describe('polling behavior', () => {
it('should clean up polling on unmount', async () => {
const mockStopPolling = vi.fn();
vi.mocked(ollamaDetection.pollForOllama).mockReturnValue(mockStopPolling);
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: false,
host: 'http://127.0.0.1:11434',
});
const { unmount } = render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText('Install Ollama'));
});
expect(ollamaDetection.pollForOllama).toHaveBeenCalled();
unmount();
expect(mockStopPolling).toHaveBeenCalled();
});
it('should handle Ollama detection during polling', async () => {
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: false,
host: 'http://127.0.0.1:11434',
});
let pollCallback: ((status: { isRunning: boolean; host: string }) => void) | undefined;
vi.mocked(ollamaDetection.pollForOllama).mockImplementation((onDetected) => {
pollCallback = onDetected;
return vi.fn();
});
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
fireEvent.click(screen.getByText('Install Ollama'));
});
expect(screen.getByText(/Waiting for Ollama/)).toBeInTheDocument();
// Simulate Ollama being detected
vi.mocked(ollamaDetection.hasModel).mockResolvedValue(true);
pollCallback!({ isRunning: true, host: 'http://127.0.0.1:11434' });
await waitFor(() => {
expect(screen.getByText('✓ Ollama is running on your system')).toBeInTheDocument();
});
});
});
describe('error states', () => {
it('should handle errors during initial check', async () => {
// Mock checkOllamaStatus to resolve with isRunning: false after an error
vi.mocked(ollamaDetection.checkOllamaStatus).mockResolvedValue({
isRunning: false,
host: 'http://127.0.0.1:11434',
error: 'Network error',
});
render(<OllamaSetup onSuccess={mockOnSuccess} onCancel={mockOnCancel} />);
await waitFor(() => {
// Should still show not detected state
expect(screen.getByText('Ollama is not detected on your system')).toBeInTheDocument();
});
});
});
});
+260
View File
@@ -0,0 +1,260 @@
import { useState, useEffect, useRef } from 'react';
import { useConfig } from './ConfigContext';
import {
checkOllamaStatus,
getOllamaDownloadUrl,
pollForOllama,
hasModel,
pullOllamaModel,
getPreferredModel,
type PullProgress,
} from '../utils/ollamaDetection';
import { initializeSystem } from '../utils/providerUtils';
import { toastService } from '../toasts';
interface OllamaSetupProps {
onSuccess: () => void;
onCancel: () => void;
}
export function OllamaSetup({ onSuccess, onCancel }: OllamaSetupProps) {
const { addExtension, getExtensions, upsert } = useConfig();
const [isChecking, setIsChecking] = useState(true);
const [ollamaDetected, setOllamaDetected] = useState(false);
const [isPolling, setIsPolling] = useState(false);
const [isConnecting, setIsConnecting] = useState(false);
const [modelStatus, setModelStatus] = useState<
'checking' | 'available' | 'not-available' | 'downloading'
>('checking');
const [downloadProgress, setDownloadProgress] = useState<PullProgress | null>(null);
const stopPollingRef = useRef<(() => void) | null>(null);
useEffect(() => {
// Check if Ollama is already running
const checkInitial = async () => {
const status = await checkOllamaStatus();
setOllamaDetected(status.isRunning);
// If Ollama is running, check for the preferred model
if (status.isRunning) {
const modelAvailable = await hasModel(getPreferredModel());
setModelStatus(modelAvailable ? 'available' : 'not-available');
}
setIsChecking(false);
};
checkInitial();
// Cleanup polling on unmount
return () => {
if (stopPollingRef.current) {
stopPollingRef.current();
}
};
}, []);
const handleInstallClick = () => {
setIsPolling(true);
// Start polling for Ollama
stopPollingRef.current = pollForOllama(
async (status) => {
setOllamaDetected(status.isRunning);
setIsPolling(false);
// Check for the model
const modelAvailable = await hasModel(getPreferredModel());
setModelStatus(modelAvailable ? 'available' : 'not-available');
toastService.success({
title: 'Ollama Detected!',
msg: 'Ollama is now running. You can connect to it.',
});
},
3000 // Check every 3 seconds
);
};
const handleDownloadModel = async () => {
setModelStatus('downloading');
setDownloadProgress({ status: 'Starting download...' });
const success = await pullOllamaModel(getPreferredModel(), (progress) => {
setDownloadProgress(progress);
});
if (success) {
setModelStatus('available');
toastService.success({
title: 'Model Downloaded!',
msg: `Successfully downloaded ${getPreferredModel()}`,
});
} else {
setModelStatus('not-available');
toastService.error({
title: 'Download Failed',
msg: `Failed to download ${getPreferredModel()}. Please try again.`,
traceback: '',
});
}
setDownloadProgress(null);
};
const handleConnectOllama = async () => {
setIsConnecting(true);
try {
// Set up Ollama configuration
await upsert('GOOSE_PROVIDER', 'ollama', false);
await upsert('GOOSE_MODEL', getPreferredModel(), false);
await upsert('OLLAMA_HOST', 'localhost', false);
// Initialize the system with Ollama
await initializeSystem('ollama', getPreferredModel(), {
getExtensions,
addExtension,
});
toastService.success({
title: 'Success!',
msg: `Connected to Ollama with ${getPreferredModel()} model.`,
});
onSuccess();
} catch (error) {
console.error('Failed to connect to Ollama:', error);
toastService.error({
title: 'Connection Failed',
msg: `Failed to connect to Ollama: ${error instanceof Error ? error.message : String(error)}`,
traceback: error instanceof Error ? error.stack || '' : '',
});
setIsConnecting(false);
}
};
if (isChecking) {
return (
<div className="space-y-4">
<div className="flex items-center justify-center py-8">
<div className="animate-spin rounded-full h-8 w-8 border-t-2 border-b-2 border-textStandard"></div>
</div>
<p className="text-center text-text-muted">Checking for Ollama...</p>
</div>
);
}
return (
<div className="space-y-6">
<div className="text-center">
<h3 className="text-lg font-semibold text-text-standard mb-2">Ollama Setup</h3>
<p className="text-text-muted">
Ollama lets you run AI models for free, private and locally on your computer.
</p>
</div>
{ollamaDetected ? (
<div className="space-y-4">
<div className="bg-background-success/10 border border-border-success rounded-lg p-4">
<p className="text-text-success text-center"> Ollama is running on your system</p>
</div>
{modelStatus === 'checking' ? (
<div className="flex items-center justify-center py-4">
<div className="animate-spin rounded-full h-6 w-6 border-t-2 border-b-2 border-textStandard"></div>
</div>
) : modelStatus === 'not-available' ? (
<div className="space-y-4">
<div className="bg-background-warning/10 border border-border-warning rounded-lg p-4">
<p className="text-text-warning text-center text-sm">
The {getPreferredModel()} model is not installed
</p>
<p className="text-text-muted text-center text-xs mt-1">
This model is recommended for the best experience with Goose
</p>
</div>
<button
onClick={handleDownloadModel}
disabled={false}
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium flex items-center justify-center gap-2"
>
Download {getPreferredModel()} (~11GB)
</button>
</div>
) : modelStatus === 'downloading' ? (
<div className="space-y-4">
<div className="bg-background-info/10 border border-border-info rounded-lg p-4">
<p className="text-text-info text-center text-sm">
Downloading {getPreferredModel()}...
</p>
{downloadProgress && (
<>
<p className="text-text-muted text-center text-xs mt-2">
{downloadProgress.status}
</p>
{downloadProgress.total && downloadProgress.completed && (
<div className="mt-3">
<div className="bg-background-muted rounded-full h-2 overflow-hidden">
<div
className="bg-background-primary h-full transition-all duration-300"
style={{
width: `${(downloadProgress.completed / downloadProgress.total) * 100}%`,
}}
/>
</div>
<p className="text-text-muted text-center text-xs mt-1">
{Math.round((downloadProgress.completed / downloadProgress.total) * 100)}%
</p>
</div>
)}
</>
)}
</div>
</div>
) : (
<button
onClick={handleConnectOllama}
disabled={isConnecting}
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium flex items-center justify-center gap-2"
>
{isConnecting ? 'Connecting...' : 'Use Goose with Ollama'}
</button>
)}
</div>
) : (
<div className="space-y-4">
<div className="bg-background-warning/10 border border-border-warning rounded-lg p-4">
<p className="text-text-warning text-center">Ollama is not detected on your system</p>
</div>
{isPolling ? (
<div className="space-y-4">
<div className="flex items-center justify-center py-4">
<div className="animate-spin rounded-full h-6 w-6 border-t-2 border-b-2 border-textStandard"></div>
</div>
<p className="text-center text-text-muted text-sm">Waiting for Ollama to start...</p>
<p className="text-center text-text-muted text-xs">
Once Ollama is installed and running, we'll automatically detect it.
</p>
</div>
) : (
<a
href={getOllamaDownloadUrl()}
target="_blank"
rel="noopener noreferrer"
onClick={handleInstallClick}
className="block w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium text-center"
>
Install Ollama
</a>
)}
</div>
)}
<button
onClick={onCancel}
className="w-full px-6 py-3 bg-transparent text-text-muted rounded-lg hover:bg-background-muted transition-colors"
>
Use a different provider
</button>
</div>
);
}
+71 -5
View File
@@ -6,6 +6,8 @@ import { startOpenRouterSetup } from '../utils/openRouterSetup';
import WelcomeGooseLogo from './WelcomeGooseLogo';
import { initializeSystem } from '../utils/providerUtils';
import { toastService } from '../toasts';
import { OllamaSetup } from './OllamaSetup';
import { checkOllamaStatus } from '../utils/ollamaDetection';
interface ProviderGuardProps {
children: React.ReactNode;
@@ -17,6 +19,8 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
const [isChecking, setIsChecking] = useState(true);
const [hasProvider, setHasProvider] = useState(false);
const [showFirstTimeSetup, setShowFirstTimeSetup] = useState(false);
const [showOllamaSetup, setShowOllamaSetup] = useState(false);
const [ollamaDetected, setOllamaDetected] = useState(false);
const [openRouterSetupState, setOpenRouterSetupState] = useState<{
show: boolean;
title: string;
@@ -101,11 +105,15 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
const provider = (await read('GOOSE_PROVIDER', false)) ?? config.GOOSE_DEFAULT_PROVIDER;
const model = (await read('GOOSE_MODEL', false)) ?? config.GOOSE_DEFAULT_MODEL;
// Always check for Ollama regardless of provider status
const ollamaStatus = await checkOllamaStatus();
setOllamaDetected(ollamaStatus.isRunning);
if (provider && model) {
console.log('ProviderGuard - Provider and model found, continuing normally');
setHasProvider(true);
} else {
console.log('ProviderGuard - No provider/model configured, showing first time setup');
console.log('ProviderGuard - No provider/model configured');
setShowFirstTimeSetup(true);
}
} catch (error) {
@@ -121,7 +129,24 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [read]);
if (isChecking && !openRouterSetupState?.show && !showFirstTimeSetup) {
// Poll for Ollama status while the first time setup is shown
useEffect(() => {
if (!showFirstTimeSetup) return;
const checkOllama = async () => {
const status = await checkOllamaStatus();
setOllamaDetected(status.isRunning);
};
// Check every 3 seconds
const interval = window.setInterval(checkOllama, 3000);
return () => {
window.clearInterval(interval);
};
}, [showFirstTimeSetup]);
if (isChecking && !openRouterSetupState?.show && !showFirstTimeSetup && !showOllamaSetup) {
return (
<div className="flex justify-center items-center py-12">
<div className="animate-spin rounded-full h-8 w-8 border-t-2 border-b-2 border-textStandard"></div>
@@ -143,6 +168,28 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
);
}
if (showOllamaSetup) {
return (
<div className="h-screen w-full flex flex-col items-center justify-center bg-background-default">
<div className="max-w-md w-full mx-auto p-8">
<div className="mb-8 text-center">
<WelcomeGooseLogo />
</div>
<OllamaSetup
onSuccess={() => {
setShowOllamaSetup(false);
setHasProvider(true);
}}
onCancel={() => {
setShowOllamaSetup(false);
setShowFirstTimeSetup(true);
}}
/>
</div>
</div>
);
}
if (showFirstTimeSetup) {
return (
<div className="h-screen w-full flex flex-col items-center justify-center bg-background-default">
@@ -156,11 +203,28 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
<div className="space-y-4">
<button
onClick={handleOpenRouterSetup}
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium"
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium flex items-center justify-center gap-2"
>
Automatic setup with OpenRouter (recommended)
</button>
<button
onClick={() => {
setShowFirstTimeSetup(false);
setShowOllamaSetup(true);
}}
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium flex items-center justify-center gap-2"
>
{ollamaDetected ? (
<>
<span className="text-text-success">●</span>
Use Ollama (auto detected)
</>
) : (
'Set up Ollama (run AI locally and free)'
)}
</button>
<button
onClick={() => navigate('/welcome', { replace: true })}
className="w-full px-6 py-3 bg-background-muted text-text-standard rounded-lg hover:bg-background-hover transition-colors font-medium"
@@ -170,8 +234,10 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
</div>
<p className="text-sm text-text-muted mt-6">
OpenRouter provides access to multiple AI models. To use this it will need to create an
account with OpenRouter.
OpenRouter provides instant access to multiple AI models with a simple setup.
{ollamaDetected
? ' Ollama is also detected on your system for running models locally.'
: ' You can also install Ollama to run free AI models locally on your computer.'}
</p>
</div>
</div>