Skip to content

Commit 78112b2

Browse files
colin-grant-workcolin-grant-work
authored
Improve image variable semantics and behavior (#16902)
- Distinguish message & context variables - Give images a place in text - Allow context from clipboard
1 parent 57530f2 commit 78112b2

16 files changed

+958
-446
lines changed

‎packages/ai-chat-ui/src/browser/chat-input-widget.tsx‎

Lines changed: 152 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import { IModelDeltaDecoration } from '@theia/monaco-editor-core/esm/vs/editor/c
4040
import { EditorOption } from '@theia/monaco-editor-core/esm/vs/editor/common/config/editorOptions';
4141
import { ChatInputHistoryService, ChatInputNavigationState } from './chat-input-history';
4242
import { ContextFileValidationService, FileValidationResult, FileValidationState } from '@theia/ai-chat/lib/browser/context-file-validation-service';
43+
import { PendingImageRegistry } from '@theia/ai-chat/lib/browser/pending-image-registry';
4344

4445
type Query = (query: string, mode?: string) => Promise<void>;
4546
type Unpin = () => void;
@@ -110,6 +111,9 @@ export class AIChatInputWidget extends ReactWidget {
110111
@inject(ContextFileValidationService) @optional()
111112
protected readonly validationService: ContextFileValidationService | undefined;
112113

114+
@inject(PendingImageRegistry)
115+
protected readonly pendingImageRegistry: PendingImageRegistry;
116+
113117
protected fileValidationState = new Map<string, FileValidationResult>();
114118

115119
@inject(ContextKeyService)
@@ -220,9 +224,19 @@ export class AIChatInputWidget extends ReactWidget {
220224

221225
protected onDisposeForChatModel = new DisposableCollection();
222226
protected _chatModel: ChatModel;
227+
228+
/**
229+
* Disposables for pending image registrations, keyed by short ID.
230+
* These are cleared after the message is sent.
231+
*/
232+
protected pendingImageDisposables = new Map<string, Disposable>();
223233
set chatModel(chatModel: ChatModel) {
224234
this.onDisposeForChatModel.dispose();
225235
this.onDisposeForChatModel = new DisposableCollection();
236+
// Register mapping from editor URI to model ID for hover provider lookups
237+
this.onDisposeForChatModel.push(
238+
this.pendingImageRegistry.registerEditorMapping(this.getResourceUri().toString(), chatModel.id)
239+
);
226240
this.onDisposeForChatModel.push(chatModel.onDidChange(event => {
227241
if (event.kind === 'addVariable') {
228242
// Validate files added via any path (including LLM tool calls)
@@ -242,21 +256,7 @@ export class AIChatInputWidget extends ReactWidget {
242256
}
243257
});
244258
this.update();
245-
} else if (event.kind === 'addRequest') {
246-
// Only clear image context variables, preserve other context (e.g., attached files)
247-
// Never clear on parse failure.
248-
const variables = chatModel.context.getVariables();
249-
const imageIndices = variables
250-
.map((v, i) => {
251-
const origin = ImageContextVariable.getOriginSafe(v);
252-
return origin === 'temporary' ? i : -1;
253-
})
254-
.filter(i => i !== -1);
255-
if (imageIndices.length > 0) {
256-
chatModel.context.deleteVariables(...imageIndices);
257-
}
258-
this.update();
259-
} else if (event.kind === 'removeVariable' || event.kind === 'changeHierarchyBranch') {
259+
} else if (event.kind === 'removeVariable' || event.kind === 'addRequest' || event.kind === 'changeHierarchyBranch') {
260260
this.update();
261261
}
262262
}));
@@ -561,38 +561,71 @@ export class AIChatInputWidget extends ReactWidget {
561561
const dataTransferText = event.dataTransfer?.getData('text/plain');
562562
const position = this.editorRef?.getControl().getTargetAtClientPoint(event.clientX, event.clientY)?.position;
563563
this.variableService.getDropResult(event.nativeEvent, { type: 'ai-chat-input-widget' }).then(result => {
564-
result.variables.forEach(variable => this.addContext(variable));
565-
const text = result.text ?? dataTransferText;
566-
if (position && text) {
564+
const textsToInsert: string[] = [];
565+
566+
result.variables.forEach(variable => {
567+
if (ImageContextVariable.isImageContextRequest(variable)) {
568+
// Register with short ID and insert short reference at drop position
569+
const shortId = this.registerPendingImage(variable);
570+
textsToInsert.push(`#${variable.variable.name}:${shortId}`);
571+
} else {
572+
this.addContext(variable);
573+
}
574+
});
575+
576+
// Add any text from drop result or data transfer
577+
if (result.text) {
578+
textsToInsert.push(result.text);
579+
} else if (dataTransferText && textsToInsert.length === 0) {
580+
// Only use dataTransferText if we have nothing else to insert
581+
textsToInsert.push(dataTransferText);
582+
}
583+
584+
if (position && textsToInsert.length > 0) {
567585
this.editorRef?.getControl().executeEdits('drag-and-drop', [{
568586
range: {
569587
startLineNumber: position.lineNumber,
570588
startColumn: position.column,
571589
endLineNumber: position.lineNumber,
572590
endColumn: position.column
573591
},
574-
text
592+
text: textsToInsert.join(' ')
575593
}]);
576594
}
577595
});
578596
}
579597

580598
protected onPaste(event: ClipboardEvent): void {
581599
this.variableService.getPasteResult(event, { type: 'ai-chat-input-widget' }).then(result => {
582-
result.variables.forEach(variable => this.addContext(variable));
583-
if (result.text) {
584-
const position = this.editorRef?.getControl().getPosition();
585-
if (position && result.text) {
586-
this.editorRef?.getControl().executeEdits('paste', [{
587-
range: {
588-
startLineNumber: position.lineNumber,
589-
startColumn: position.column,
590-
endLineNumber: position.lineNumber,
591-
endColumn: position.column
592-
},
593-
text: result.text
594-
}]);
600+
const position = this.editorRef?.getControl().getPosition();
601+
const textsToInsert: string[] = [];
602+
603+
result.variables.forEach(variable => {
604+
if (ImageContextVariable.isImageContextRequest(variable)) {
605+
// Register with short ID and insert short reference at cursor position
606+
const shortId = this.registerPendingImage(variable);
607+
textsToInsert.push(`#${variable.variable.name}:${shortId}`);
608+
} else {
609+
this.addContext(variable);
595610
}
611+
});
612+
613+
// Insert any text from the paste result
614+
if (result.text) {
615+
textsToInsert.push(result.text);
616+
}
617+
618+
// Insert all collected text at cursor position
619+
if (position && textsToInsert.length > 0) {
620+
this.editorRef?.getControl().executeEdits('paste', [{
621+
range: {
622+
startLineNumber: position.lineNumber,
623+
startColumn: position.column,
624+
endLineNumber: position.lineNumber,
625+
endColumn: position.column
626+
},
627+
text: textsToInsert.join(' ')
628+
}]);
596629
}
597630
});
598631
}
@@ -631,10 +664,6 @@ export class AIChatInputWidget extends ReactWidget {
631664
});
632665
}
633666

634-
protected deleteContextElement(index: number): void {
635-
this._chatModel.context.deleteVariables(index);
636-
}
637-
638667
protected handleContextMenu(event: IMouseEvent): void {
639668
this.contextMenuRenderer.render({
640669
menuPath: AIChatInputWidget.CONTEXT_MENU,
@@ -650,9 +679,87 @@ export class AIChatInputWidget extends ReactWidget {
650679
this._chatModel.context.addVariables(variable);
651680
}
652681

682+
/**
683+
* Get the scope URI used for registering pending images.
684+
* Uses the chat model ID for scoping, not the editor URI.
685+
*/
686+
protected getScopeUri(): string {
687+
const modelId = this._chatModel?.id ?? 'default';
688+
return this.pendingImageRegistry.getScopeUriForModel(modelId);
689+
}
690+
691+
/**
692+
* Register a pending image attachment with a short ID.
693+
* @returns The short ID that can be used in the text.
694+
*/
695+
registerPendingImage(variable: AIVariableResolutionRequest): string {
696+
const parsed = ImageContextVariable.parseRequest(variable);
697+
if (!parsed || !variable.arg) {
698+
throw new Error('Invalid image context variable');
699+
}
700+
701+
const scopeUri = this.getScopeUri();
702+
const baseName = this.getImageBaseName(parsed);
703+
const shortId = this.pendingImageRegistry.generateShortId(baseName, scopeUri);
704+
705+
const disposable = this.pendingImageRegistry.register(
706+
scopeUri,
707+
shortId,
708+
parsed,
709+
variable.arg
710+
);
711+
this.pendingImageDisposables.set(shortId, disposable);
712+
713+
return shortId;
714+
}
715+
716+
/**
717+
* Clear all pending image attachments. Called after a message is sent.
718+
*/
719+
clearPendingImageAttachments(): void {
720+
// Dispose all registrations
721+
for (const disposable of this.pendingImageDisposables.values()) {
722+
disposable.dispose();
723+
}
724+
this.pendingImageDisposables.clear();
725+
726+
// Also clear the scope in the registry (belt and suspenders)
727+
this.pendingImageRegistry.clearScope(this.getScopeUri());
728+
}
729+
730+
/**
731+
* Get a meaningful base name for an image variable.
732+
* For dropped files (with wsRelativePath), use the workspace-relative path (like file variables).
733+
* For pasted images, use "pasted_image".
734+
*/
735+
protected getImageBaseName(imageVariable: ImageContextVariable): string {
736+
// For file-based images, use the workspace-relative path (consistent with #file: variables)
737+
if (imageVariable.wsRelativePath) {
738+
return imageVariable.wsRelativePath;
739+
}
740+
741+
// For pasted images, use a generic name
742+
return 'pasted_image';
743+
}
744+
745+
/**
746+
* Get all variables to be sent with the request.
747+
* Note: Image variables in text use short IDs that get resolved via the pending image registry.
748+
*/
749+
getAllVariablesForRequest(): AIVariableResolutionRequest[] {
750+
return [...this._chatModel.context.getVariables()];
751+
}
752+
653753
protected getContext(): readonly AIVariableResolutionRequest[] {
654754
return this._chatModel.context.getVariables();
655755
}
756+
757+
/**
758+
* Delete a context element by index.
759+
*/
760+
protected deleteContextElement(index: number): void {
761+
this._chatModel.context.deleteVariables(index);
762+
}
656763
}
657764

658765
interface ChatInputProperties {
@@ -776,10 +883,10 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
776883

777884
const editor = await props.editorProvider.createSimpleInline(uri, editorContainerRef.current!, {
778885
language: CHAT_VIEW_LANGUAGE_EXTENSION,
779-
// Disable code lens, inlay hints and hover support to avoid console errors from other contributions
886+
// Disable code lens and inlay hints to avoid console errors from other contributions
780887
codeLens: false,
781888
inlayHints: { enabled: 'off' },
782-
hover: { enabled: false },
889+
hover: { enabled: true },
783890
autoSizing: false, // we handle the sizing ourselves
784891
scrollBeyondLastLine: false,
785892
scrollBeyondLastColumn: 0,
@@ -1475,33 +1582,29 @@ const ChatContext: React.FunctionComponent<ChatContextUI> = ({ context }) => (
14751582
<ul>
14761583
{context.map((element, index) => {
14771584
if (ImageContextVariable.isImageContextRequest(element.variable)) {
1478-
let variable: ImageContextVariable | undefined;
1479-
try {
1480-
variable = ImageContextVariable.parseRequest(element.variable);
1481-
} catch {
1482-
variable = undefined;
1483-
}
1484-
1585+
const variable = ImageContextVariable.parseRequest(element.variable);
1586+
const displayName = variable?.name ?? variable?.wsRelativePath?.split('/').pop() ?? element.name;
14851587
const title = variable?.name ?? variable?.wsRelativePath ?? element.details ?? element.name;
1486-
const label = variable?.name ?? variable?.wsRelativePath?.split('/').pop() ?? element.name;
1588+
// Check if we have the actual image data (pre-processed) or just a path reference
1589+
const hasImageData = variable && ImageContextVariable.isResolved(variable);
14871590

14881591
return <li key={index} className="theia-ChatInput-ChatContext-Element theia-ChatInput-ImageContext-Element"
14891592
title={title} onClick={() => element.open?.()}>
14901593
<div className="theia-ChatInput-ChatContext-Row">
14911594
<div className={`theia-ChatInput-ChatContext-Icon ${element.iconClass}`} />
14921595
<div className="theia-ChatInput-ChatContext-labelParts">
14931596
<span className={`theia-ChatInput-ChatContext-title ${element.nameClass}`}>
1494-
{label}
1597+
{displayName}
14951598
</span>
14961599
<span className='theia-ChatInput-ChatContext-additionalInfo'>
14971600
{element.additionalInfo}
14981601
</span>
14991602
</div>
15001603
<span className="codicon codicon-close action" title={nls.localizeByDefault('Delete')} onClick={e => { e.stopPropagation(); element.delete(); }} />
15011604
</div>
1502-
{variable && <div className="theia-ChatInput-ChatContext-ImageRow">
1605+
{hasImageData && <div className="theia-ChatInput-ChatContext-ImageRow">
15031606
<div className='theia-ChatInput-ImagePreview-Item'>
1504-
<img src={`data:${variable.mimeType};base64,${variable.data}`} alt={variable.name ?? label} />
1607+
<img src={`data:${variable.mimeType};base64,${variable.data}`} alt={variable.name ?? displayName} />
15051608
</div>
15061609
</div>}
15071610
</li>;

‎packages/ai-chat-ui/src/browser/chat-tree-view/chat-view-tree-input-widget.tsx‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,10 +94,29 @@ export class AIChatTreeInputWidget extends AIChatInputWidget {
9494
this.request.editContextManager.addVariables(variable);
9595
}
9696

97+
/**
98+
* In edit mode, image attachments go to editContextManager like other context.
99+
* They don't need separate handling since editContextManager is message-scoped.
100+
* We add the variable directly rather than registering with the pending image registry.
101+
*/
102+
override registerPendingImage(variable: AIVariableResolutionRequest): string {
103+
this.request.editContextManager.addVariables(variable);
104+
// Return a placeholder short ID - in edit mode, the full data is in editContextManager
105+
return `edit_img_${Date.now()}`;
106+
}
107+
97108
protected override getContext(): readonly AIVariableResolutionRequest[] {
98109
return this.request.editContextManager.getVariables();
99110
}
100111

112+
override getAllVariablesForRequest(): AIVariableResolutionRequest[] {
113+
return [...this.request.editContextManager.getVariables()];
114+
}
115+
116+
override clearPendingImageAttachments(): void {
117+
// No-op in edit mode - editContextManager handles its own lifecycle
118+
}
119+
101120
protected override deleteContextElement(index: number): void {
102121
this.request.editContextManager.deleteVariables(index);
103122
}

‎packages/ai-chat-ui/src/browser/chat-tree-view/chat-view-tree-widget.tsx‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ import {
2828
type ChatRequest,
2929
type ChatHierarchyBranch,
3030
} from '@theia/ai-chat';
31-
import { ImageContextVariable } from '@theia/ai-chat/lib/common/image-context-variable';
3231
import { AIVariableService } from '@theia/ai-core';
3332
import { AIActivationService } from '@theia/ai-core/lib/browser';
3433
import { CommandRegistry, ContributionProvider, Disposable, DisposableCollection, Emitter, Event } from '@theia/core';
@@ -830,31 +829,6 @@ const ChatRequestRender = (
830829
);
831830
};
832831

833-
// Extract image variables from the request context
834-
const imageVariables = node.request.context.variables
835-
.filter(ImageContextVariable.isResolvedImageContext)
836-
.map(resolved => ImageContextVariable.parseResolved(resolved))
837-
.filter((img): img is NonNullable<typeof img> => img !== undefined);
838-
839-
const renderImages = () => {
840-
if (imageVariables.length === 0) {
841-
return undefined;
842-
}
843-
return (
844-
<div className="theia-RequestNode-Images">
845-
{imageVariables.map((img, index) => (
846-
<div key={index} className="theia-RequestNode-ImagePreview">
847-
<img
848-
src={`data:${img.mimeType};base64,${img.data}`}
849-
alt={img.name ?? img.wsRelativePath ?? 'Image'}
850-
title={img.name ?? img.wsRelativePath ?? 'Image'}
851-
/>
852-
</div>
853-
))}
854-
</div>
855-
);
856-
};
857-
858832
return (
859833
<div className="theia-RequestNode">
860834
<p>
@@ -895,7 +869,6 @@ const ChatRequestRender = (
895869
}
896870
})}
897871
</p>
898-
{renderImages()}
899872
{renderFooter()}
900873
</div>
901874
);

0 commit comments

Comments
 (0)