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
29 changes: 28 additions & 1 deletion src/components/AIView.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import { getModelName } from "@/utils/ai/models";
import type { DisplayedMessage } from "@/types/message";
import type { RuntimeConfig } from "@/types/runtime";
import { useAIViewKeybinds } from "@/hooks/useAIViewKeybinds";
import { evictModelFromLRU } from "@/hooks/useModelLRU";

interface AIViewProps {
workspaceId: string;
Expand Down Expand Up @@ -66,9 +67,35 @@ const AIViewInner: React.FC<AIViewProps> = ({
storageKey: "review-sidebar-width", // Persists across sessions
});

// Get workspace state from store (only re-renders when THIS workspace changes)
const workspaceState = useWorkspaceState(workspaceId);
const aggregator = useWorkspaceAggregator(workspaceId);
const handledModelErrorsRef = useRef<Set<string>>(new Set());

useEffect(() => {
handledModelErrorsRef.current.clear();
}, [workspaceId]);

useEffect(() => {
if (!workspaceState) {
return;
}

for (const message of workspaceState.messages) {
if (message.type !== "stream-error") {
continue;
}
if (message.errorType !== "model_not_found") {
continue;
}
if (handledModelErrorsRef.current.has(message.id)) {
continue;
}
handledModelErrorsRef.current.add(message.id);
if (message.model) {
evictModelFromLRU(message.model);
}
}
}, [workspaceState, workspaceId]);

const [editingMessage, setEditingMessage] = useState<{ id: string; content: string } | undefined>(
undefined
Expand Down
3 changes: 2 additions & 1 deletion src/components/ChatInput/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
const inputRef = useRef<HTMLTextAreaElement>(null);
const modelSelectorRef = useRef<ModelSelectorRef>(null);
const [mode, setMode] = useMode();
const { recentModels, addModel } = useModelLRU();
const { recentModels, addModel, evictModel } = useModelLRU();
const commandListId = useId();
const telemetry = useTelemetry();
const [vimEnabled, setVimEnabled] = usePersistedState<boolean>(VIM_ENABLED_KEY, false, {
Expand Down Expand Up @@ -922,6 +922,7 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
value={preferredModel}
onChange={setPreferredModel}
recentModels={recentModels}
onRemoveModel={evictModel}
onComplete={() => inputRef.current?.focus()}
/>
<TooltipWrapper inline>
Expand Down
7 changes: 7 additions & 0 deletions src/components/ModelSelector.stories.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ const meta = {
control: false,
description: "Callback when model changes",
},
onRemoveModel: {
control: false,
description: "Callback when a model is removed",
},
recentModels: {
control: { type: "object" },
description: "List of recently used models",
Expand All @@ -36,6 +40,7 @@ export const Default: Story = {
args: {
value: "anthropic:claude-sonnet-4-5",
onChange: action("onChange"),
onRemoveModel: action("onRemoveModel"),
recentModels: ["anthropic:claude-sonnet-4-5", "anthropic:claude-opus-4-1", "openai:gpt-5-pro"],
onComplete: action("onComplete"),
},
Expand All @@ -45,6 +50,7 @@ export const LongModelName: Story = {
args: {
value: "anthropic:claude-opus-4-20250514-preview-experimental",
onChange: action("onChange"),
onRemoveModel: action("onRemoveModel"),
recentModels: [
"anthropic:claude-opus-4-20250514-preview-experimental",
"anthropic:claude-sonnet-4-20250514-preview-experimental",
Expand All @@ -58,6 +64,7 @@ export const WithManyModels: Story = {
args: {
value: "anthropic:claude-sonnet-4-5",
onChange: action("onChange"),
onRemoveModel: action("onRemoveModel"),
recentModels: [
"anthropic:claude-sonnet-4-5",
"anthropic:claude-opus-4-1",
Expand Down
33 changes: 31 additions & 2 deletions src/components/ModelSelector.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ interface ModelSelectorProps {
value: string;
onChange: (value: string) => void;
recentModels: string[];
onRemoveModel?: (model: string) => void;
onComplete?: () => void;
}

Expand All @@ -20,7 +21,7 @@ export interface ModelSelectorRef {
}

export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
({ value, onChange, recentModels, onComplete }, ref) => {
({ value, onChange, recentModels, onRemoveModel, onComplete }, ref) => {
const [isEditing, setIsEditing] = useState(false);
const [inputValue, setInputValue] = useState(value);
const [error, setError] = useState<string | null>(null);
Expand Down Expand Up @@ -151,6 +152,22 @@ export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
setShowDropdown(false);
};

const handleRemoveModel = useCallback(
(model: string, event: React.MouseEvent<HTMLButtonElement>) => {
event.preventDefault();
event.stopPropagation();
if (!onRemoveModel) {
return;
}
onRemoveModel(model);
setHighlightedIndex(-1);
if (inputValue === model) {
setInputValue("");
}
},
[inputValue, onRemoveModel]
);

const handleClick = useCallback(() => {
setIsEditing(true);
setInputValue(""); // Clear input to show all models
Expand Down Expand Up @@ -224,7 +241,19 @@ export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
)}
onClick={() => handleSelectModel(model)}
>
{model}
<div className="flex items-center justify-between gap-2">
<span className="truncate">{model}</span>
{onRemoveModel && (
<button
type="button"
onClick={(event) => handleRemoveModel(model, event)}
className="text-muted-light border-border-light/40 hover:border-danger-soft/60 hover:text-danger-soft rounded-sm border px-1 py-0.5 text-[9px] font-semibold tracking-wide uppercase transition-colors duration-150"
aria-label={`Remove ${model} from recent models`}
>
×
</button>
)}
</div>
</div>
))}
</div>
Expand Down
30 changes: 28 additions & 2 deletions src/hooks/useModelLRU.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { useCallback, useEffect } from "react";
import { usePersistedState, readPersistedState } from "./usePersistedState";
import { usePersistedState, readPersistedState, updatePersistedState } from "./usePersistedState";
import { MODEL_ABBREVIATIONS } from "@/utils/slashCommands/registry";
import { defaultModel } from "@/utils/ai/models";

Expand All @@ -12,6 +12,23 @@ const DEFAULT_MODELS = [
defaultModel,
...Object.values(MODEL_ABBREVIATIONS).filter((m) => m !== defaultModel),
].slice(0, MAX_LRU_SIZE);
function persistModels(models: string[]): void {
updatePersistedState(LRU_KEY, models.slice(0, MAX_LRU_SIZE));
}

export function evictModelFromLRU(model: string): void {
const normalized = model.trim();
if (!normalized) {
return;
}
const current = readPersistedState<string[]>(LRU_KEY, DEFAULT_MODELS.slice(0, MAX_LRU_SIZE));
const filtered = current.filter((m) => m !== normalized);
if (filtered.length === current.length) {
return;
}
const nextList = filtered.length > 0 ? filtered : DEFAULT_MODELS.slice(0, MAX_LRU_SIZE);
persistModels(nextList);
}

/**
* Get the default model from LRU (non-hook version for use outside React)
Expand All @@ -32,7 +49,8 @@ export function getDefaultModelFromLRU(): string {
export function useModelLRU() {
const [recentModels, setRecentModels] = usePersistedState<string[]>(
LRU_KEY,
DEFAULT_MODELS.slice(0, MAX_LRU_SIZE)
DEFAULT_MODELS.slice(0, MAX_LRU_SIZE),
{ listener: true }
);

// Merge any new defaults from MODEL_ABBREVIATIONS (only once on mount)
Expand Down Expand Up @@ -76,8 +94,16 @@ export function useModelLRU() {
return recentModels;
}, [recentModels]);

const evictModel = useCallback((modelString: string) => {
if (!modelString.trim()) {
return;
}
evictModelFromLRU(modelString);
}, []);

return {
addModel,
evictModel,
getRecentModels,
recentModels,
};
Expand Down