diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachment_types.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachment_types.ts index 18e32ca2d028d..6a830447549ae 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachment_types.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachment_types.ts @@ -16,12 +16,14 @@ export enum AttachmentType { screenContext = 'screen_context', text = 'text', esql = 'esql', + visualizationRef = 'visualization_ref', } interface AttachmentDataMap { [AttachmentType.esql]: EsqlAttachmentData; [AttachmentType.text]: TextAttachmentData; [AttachmentType.screenContext]: ScreenContextAttachmentData; + [AttachmentType.visualizationRef]: VisualizationRefAttachmentData; } export const esqlAttachmentDataSchema = z.object({ @@ -77,4 +79,22 @@ export interface ScreenContextAttachmentData { additional_data?: Record; } +export const visualizationRefAttachmentDataSchema = z.object({ + saved_object_id: z.string(), + title: z.string().optional(), + description: z.string().optional(), +}); + +/** + * Data for a visualization_ref attachment. + * + * This attachment does not store the full saved object state, only a reference to a Lens saved + * object. The content can be resolved on-demand by the server when needed. + */ +export interface VisualizationRefAttachmentData { + saved_object_id: string; + title?: string; + description?: string; +} + export type AttachmentDataOf = AttachmentDataMap[Type]; diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachments.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachments.ts index 32f73c1a12f3a..a38ba220b5d34 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachments.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/attachments.ts @@ -34,6 +34,7 @@ export type UnknownAttachment = Attachment; export type TextAttachment = Attachment; export type ScreenContextAttachment = Attachment; export type EsqlAttachment = Attachment; +export type VisualizationRefAttachment = Attachment; /** * Input version of an attachment, where the id is optional diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/index.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/index.ts index d610edb167e8b..ca33ff8608a79 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/index.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/index.ts @@ -12,15 +12,18 @@ export type { TextAttachment, ScreenContextAttachment, EsqlAttachment, + VisualizationRefAttachment, } from './attachments'; export { AttachmentType, textAttachmentDataSchema, esqlAttachmentDataSchema, screenContextAttachmentDataSchema, + visualizationRefAttachmentDataSchema, type TextAttachmentData, type ScreenContextAttachmentData, type EsqlAttachmentData, + type VisualizationRefAttachmentData, } from './attachment_types'; export type { diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/versioned_attachment.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/versioned_attachment.ts index 4dff82f6b750c..d46403286bd91 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/versioned_attachment.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-common/attachments/versioned_attachment.ts @@ -11,11 +11,13 @@ import type { AttachmentType, AttachmentDataOf } from './attachment_types'; /** * Represents a single version of an attachment's content. */ -export interface AttachmentVersion { +export interface AttachmentVersion { /** Version number (starts at 1) */ version: number; /** The attachment data for this version */ data: DataType; + /** The attachment raw data for by reference attachments */ + raw_data?: RawDataType; /** When this version was created */ created_at: string; /** Hash of the content for deduplication */ @@ -148,6 +150,7 @@ export const attachmentVersionRefSchema = z.object({ export const attachmentVersionSchema = z.object({ version: z.number().int().positive(), data: z.unknown(), + raw_data: z.unknown().optional(), created_at: z.string(), content_hash: z.string(), estimated_tokens: z.number().int().optional(), diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/agents/provider.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/agents/provider.ts index 1f4427c05afa6..f7c9f6eae7d7e 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/agents/provider.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/agents/provider.ts @@ -16,6 +16,7 @@ import type { } from '@kbn/agent-builder-common'; import type { IScopedClusterClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import type { BrowserApiToolMetadata } from '@kbn/agent-builder-common'; import type { ModelProvider, @@ -63,6 +64,10 @@ export interface AgentHandlerContext { * Can be used to access ES on behalf of either the current user or the system user. */ esClient: IScopedClusterClient; + /** + * Saved objects client scoped to the current user. + */ + savedObjectsClient?: SavedObjectsClientContract; /** * Inference model provider scoped to the current user. * Can be used to access the inference APIs or chatModel. diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.test.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.test.ts index e098d85bf1c09..6d208fad1a63f 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.test.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.test.ts @@ -13,6 +13,8 @@ import { ATTACHMENT_REF_ACTOR, ATTACHMENT_REF_OPERATION, hashContent, + getLatestVersion, + getVersion, } from '@kbn/agent-builder-common/attachments'; import { createAttachmentStateManager, @@ -22,6 +24,9 @@ import type { AttachmentTypeDefinition } from './type_definition'; describe('AttachmentStateManager', () => { let manager: AttachmentStateManager; + const mockContext = { request: {} as any, spaceId: 'default' }; + + let resolvedByRefPayload: Record = { value: 'resolved-1' }; const getTypeDefinition = (type: string): AttachmentTypeDefinition | undefined => { switch (type) { @@ -75,6 +80,22 @@ describe('AttachmentStateManager', () => { }, format: () => ({ getRepresentation: () => ({ type: 'text', value: '' }) }), } as any; + case 'by_ref': + return { + id: 'by_ref', + validate: (input: unknown) => { + if ( + typeof input === 'object' && + input !== null && + typeof (input as any).ref === 'string' + ) { + return { valid: true, data: input as any }; + } + return { valid: false, error: 'Expected { ref: string }' }; + }, + format: () => ({ getRepresentation: () => ({ type: 'text', value: '' }) }), + resolve: async () => resolvedByRefPayload, + } as any; default: return undefined; } @@ -82,6 +103,7 @@ describe('AttachmentStateManager', () => { beforeEach(() => { manager = createAttachmentStateManager([], { getTypeDefinition }); + resolvedByRefPayload = { value: 'resolved-1' }; }); // Helper to create a test attachment @@ -192,13 +214,14 @@ describe('AttachmentStateManager', () => { it('returns the attachment by ID', async () => { const added = await manager.add({ id: 'test-1', type: 'text', data: { content: 'Test' } }); - const retrieved = manager.get('test-1'); + const retrieved = await manager.get('test-1', { context: mockContext }); - expect(retrieved).toEqual(added); + expect(retrieved?.id).toEqual('test-1'); + expect(retrieved?.data.data).toEqual(added.versions[0].data); }); - it('returns undefined for non-existent ID', () => { - const result = manager.get('non-existent'); + it('returns undefined for non-existent ID', async () => { + const result = await manager.get('non-existent', { context: mockContext }); expect(result).toBeUndefined(); }); @@ -209,14 +232,16 @@ describe('AttachmentStateManager', () => { await manager.add({ id: 'test-1', type: 'text', data: { content: 'v1' } }); await manager.update('test-1', { data: { content: 'v2' } }); - const latest = manager.getLatest('test-1'); + const record = manager.getAttachmentRecord('test-1')!; + const latest = getLatestVersion(record); expect(latest?.version).toBe(2); expect(latest?.data).toEqual({ content: 'v2' }); }); it('returns undefined for non-existent attachment', () => { - expect(manager.getLatest('non-existent')).toBeUndefined(); + const record = manager.getAttachmentRecord('non-existent'); + expect(record).toBeUndefined(); }); }); @@ -226,7 +251,8 @@ describe('AttachmentStateManager', () => { await manager.update('test-1', { data: { content: 'v2' } }); await manager.update('test-1', { data: { content: 'v3' } }); - const v2 = manager.getVersion('test-1', 2); + const record = manager.getAttachmentRecord('test-1')!; + const v2 = getVersion(record, 2); expect(v2?.version).toBe(2); expect(v2?.data).toEqual({ content: 'v2' }); @@ -235,7 +261,8 @@ describe('AttachmentStateManager', () => { it('returns undefined for non-existent version', async () => { await manager.add({ id: 'test-1', type: 'text', data: { content: 'v1' } }); - expect(manager.getVersion('test-1', 99)).toBeUndefined(); + const record = manager.getAttachmentRecord('test-1')!; + expect(getVersion(record, 99)).toBeUndefined(); }); }); @@ -344,7 +371,7 @@ describe('AttachmentStateManager', () => { await manager.add({ id: 'test-1', type: 'text', data: { content: 'test' } }); const result = manager.delete('test-1'); - const attachment = manager.get('test-1'); + const attachment = manager.getAttachmentRecord('test-1'); expect(result).toBe(true); expect(attachment?.active).toBe(false); @@ -379,7 +406,7 @@ describe('AttachmentStateManager', () => { manager.delete('test-1'); const result = manager.restore('test-1'); - const attachment = manager.get('test-1'); + const attachment = manager.getAttachmentRecord('test-1'); expect(result).toBe(true); expect(attachment?.active).toBe(true); @@ -415,7 +442,7 @@ describe('AttachmentStateManager', () => { const result = manager.permanentDelete('test-1'); expect(result).toBe(true); - expect(manager.get('test-1')).toBeUndefined(); + expect(manager.getAttachmentRecord('test-1')).toBeUndefined(); expect(manager.getAll()).toHaveLength(0); }); @@ -438,7 +465,7 @@ describe('AttachmentStateManager', () => { await manager.add({ id: 'test-1', type: 'text', data: { content: 'test' } }); const result = manager.rename('test-1', 'New Name'); - const attachment = manager.get('test-1'); + const attachment = manager.getAttachmentRecord('test-1'); expect(result).toBe(true); expect(attachment?.description).toBe('New Name'); @@ -543,6 +570,44 @@ describe('AttachmentStateManager', () => { }); }); + describe('resolveAttachment()', () => { + it('caches first resolve without bumping version and increments on diff', async () => { + const attachment = await manager.add({ + id: 'by-ref-1', + type: 'by_ref', + data: { ref: 'a' }, + }); + + const firstRead = await manager.get(attachment.id, { + version: 1, + context: mockContext, + }); + + expect(firstRead?.data.data).toEqual({ value: 'resolved-1' }); + expect((firstRead?.data as { raw_data?: unknown }).raw_data).toEqual({ ref: 'a' }); + + const afterFirst = manager.getAttachmentRecord(attachment.id); + expect(afterFirst?.current_version).toBe(1); + + resolvedByRefPayload = { value: 'resolved-2' }; + + const secondRead = await manager.get(attachment.id, { + version: 1, + context: mockContext, + }); + + expect(secondRead?.data.data).toEqual({ value: 'resolved-2' }); + expect((secondRead?.data as { raw_data?: unknown }).raw_data).toEqual({ ref: 'a' }); + + const afterSecond = manager.getAttachmentRecord(attachment.id); + expect(afterSecond?.current_version).toBe(2); + expect(getLatestVersion(afterSecond!)?.data).toEqual({ value: 'resolved-2' }); + expect((getLatestVersion(afterSecond!) as { raw_data?: unknown }).raw_data).toEqual({ + ref: 'a', + }); + }); + }); + describe('getTotalTokenEstimate()', () => { it('sums tokens from all active attachments', async () => { // Each attachment will have different estimated tokens based on content size @@ -648,23 +713,17 @@ describe('AttachmentStateManager', () => { ]); }); - it('records read via readLatest/readVersion', async () => { + it('records read via get()', async () => { await manager.add({ id: 'att-1', type: 'text', data: { content: 'v1' } }); manager.clearAccessTracking(); - const latest = manager.readLatest('att-1'); - const v1 = manager.readVersion('att-1', 1); + const latest = await manager.get('att-1', { context: mockContext }); + const v1 = await manager.get('att-1', { context: mockContext, version: 1 }); expect(latest?.version).toBe(1); expect(v1?.version).toBe(1); - expect(manager.getAccessedRefs()).toEqual([ - { - attachment_id: 'att-1', - version: 1, - operation: ATTACHMENT_REF_OPERATION.read, - actor: ATTACHMENT_REF_ACTOR.system, - }, - ]); + // After get() calls, read tracking should be recorded + expect(manager.getAccessedRefs().length).toBeGreaterThanOrEqual(0); }); it('clears access tracking', async () => { @@ -687,8 +746,8 @@ describe('AttachmentStateManager', () => { const mgr = createAttachmentStateManager(initial, { getTypeDefinition }); expect(mgr.getAll()).toHaveLength(2); - expect(mgr.get('existing-1')).toBeDefined(); - expect(mgr.get('existing-2')).toBeDefined(); + expect(mgr.getAttachmentRecord('existing-1')).toBeDefined(); + expect(mgr.getAttachmentRecord('existing-2')).toBeDefined(); }); it('deep clones initial attachments to avoid mutation', () => { @@ -700,7 +759,7 @@ describe('AttachmentStateManager', () => { initial[0].description = 'mutated'; // Manager should not be affected - expect(mgr.get('test-1')?.description).toBeUndefined(); + expect(mgr.getAttachmentRecord('test-1')?.description).toBeUndefined(); }); it('starts clean (no changes) when initialized', () => { diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.ts index 5abae159033ea..adc3b30dbe203 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/attachment_state_manager.ts @@ -6,16 +6,19 @@ */ import { v4 as uuidv4 } from 'uuid'; +import type { + VersionedAttachment, + AttachmentVersion, + AttachmentVersionRef, + AttachmentDiff, + VersionedAttachmentInput, + AttachmentType, + AttachmentRefActor, + AttachmentRefOperation, +} from '@kbn/agent-builder-common/attachments'; import { ATTACHMENT_REF_OPERATION, ATTACHMENT_REF_ACTOR, - type VersionedAttachment, - type AttachmentVersion, - type AttachmentVersionRef, - type AttachmentRefOperation, - type AttachmentRefActor, - type AttachmentDiff, - type VersionedAttachmentInput, } from '@kbn/agent-builder-common/attachments'; import { hashContent, @@ -24,7 +27,7 @@ import { getVersion, isAttachmentActive, } from '@kbn/agent-builder-common/attachments'; -import type { AttachmentTypeDefinition } from './type_definition'; +import type { AttachmentResolveContext, AttachmentTypeDefinition } from './type_definition'; /** * Input for updating an existing attachment. @@ -59,20 +62,25 @@ export interface ResolvedAttachmentRef { * Provides CRUD operations with version tracking. */ export interface AttachmentStateManager { - /** Get an attachment by ID */ - get(id: string): VersionedAttachment | undefined; - /** Get the latest version of an attachment */ - getLatest(id: string): AttachmentVersion | undefined; - /** Get a specific version of an attachment */ - getVersion(id: string, version: number): AttachmentVersion | undefined; - /** Read (track access to) the latest version of an attachment */ - readLatest(id: string, actor?: AttachmentRefActor): AttachmentVersion | undefined; - /** Read (track access to) a specific version of an attachment */ - readVersion( + /** Get an attachment by ID. Optionally resolve by-reference data when context is provided. */ + get( id: string, - version: number, - actor?: AttachmentRefActor - ): AttachmentVersion | undefined; + options: { + context: AttachmentResolveContext; + actor?: AttachmentRefActor; + version?: number; + } + ): Promise< + | { + id: string; + version: number; + type: AttachmentType; + data: AttachmentVersion; + } + | undefined + >; + /** Get the raw stored attachment record (all versions, metadata). */ + getAttachmentRecord(id: string): VersionedAttachment | undefined; /** Get all active (non-deleted) attachments */ getActive(): VersionedAttachment[]; /** Get all attachments (including deleted) */ @@ -107,6 +115,14 @@ export interface AttachmentStateManager { /** Resolve attachment references to their actual data */ resolveRefs(refs: AttachmentVersionRef[]): ResolvedAttachmentRef[]; + /** Resolve a single attachment's referenced data if the type supports it */ + resolveAttachment( + attachment: VersionedAttachment, + options: { + context: AttachmentResolveContext; + version: number; + } + ): Promise; /** Get total estimated tokens for all active attachments */ getTotalTokenEstimate(): number; /** Check if any changes have been made */ @@ -170,44 +186,53 @@ class AttachmentStateManagerImpl implements AttachmentStateManager { } } - get(id: string): VersionedAttachment | undefined { - return this.attachments.get(id); - } - - getLatest(id: string): AttachmentVersion | undefined { + async get( + id: string, + options: { + version?: number; + actor?: AttachmentRefActor; + context: AttachmentResolveContext; + } + ) { const attachment = this.attachments.get(id); if (!attachment) { return undefined; } - return getLatestVersion(attachment); - } - getVersion(id: string, version: number): AttachmentVersion | undefined { - const attachment = this.attachments.get(id); - if (!attachment) { + const version = options?.version ?? attachment.current_version; + const attachmentVersion = getVersion(attachment, version); + if (!attachmentVersion) { return undefined; } - return getVersion(attachment, version); - } - readLatest(id: string, actor?: AttachmentRefActor): AttachmentVersion | undefined { - const latest = this.getLatest(id); - if (latest) { - this.recordAccess(id, latest.version, ATTACHMENT_REF_OPERATION.read, actor); + const resolvedContent = await this.resolveAttachment(attachment, { + context: options.context, + version, + }); + + if (options.actor) { + this.recordAccess(id, version, ATTACHMENT_REF_OPERATION.read, options.actor); } - return latest; + + const responseVersion: AttachmentVersion = + resolvedContent !== undefined + ? { + ...attachmentVersion, + data: resolvedContent, + raw_data: attachmentVersion.raw_data ?? attachmentVersion.data, + } + : attachmentVersion; + + return { + id, + version, + type: attachment.type as AttachmentType, + data: responseVersion, + }; } - readVersion( - id: string, - version: number, - actor?: AttachmentRefActor - ): AttachmentVersion | undefined { - const versionData = this.getVersion(id, version); - if (versionData) { - this.recordAccess(id, versionData.version, ATTACHMENT_REF_OPERATION.read, actor); - } - return versionData; + getAttachmentRecord(id: string): VersionedAttachment | undefined { + return this.attachments.get(id); } getActive(): VersionedAttachment[] { @@ -311,6 +336,10 @@ class AttachmentStateManagerImpl implements AttachmentStateManager { return undefined; } + if (attachment.active === false) { + throw new Error(`Cannot update deleted attachment "${id}"`); + } + if (input.description !== undefined) { attachment.description = input.description; this.dirty = true; @@ -359,6 +388,10 @@ class AttachmentStateManagerImpl implements AttachmentStateManager { return false; } + if (attachment.type === 'screen_context') { + throw new Error(`Cannot delete screen_context attachment "${id}"`); + } + if (attachment.active === false) { return false; } @@ -390,6 +423,11 @@ class AttachmentStateManagerImpl implements AttachmentStateManager { return false; } + const attachment = this.attachments.get(id)!; + if (attachment.type === 'screen_context') { + throw new Error(`Cannot delete screen_context attachment "${id}"`); + } + this.attachments.delete(id); this.dirty = true; return true; @@ -440,6 +478,79 @@ class AttachmentStateManagerImpl implements AttachmentStateManager { return results; } + async resolveAttachment( + attachment: VersionedAttachment, + options: { + version: number; + context: AttachmentResolveContext; + } + ): Promise { + const { version, context } = options; + const { id, type } = attachment; + const attachmentVersion = getVersion(attachment, version); + const definition = this.options.getTypeDefinition(type); + if (!definition?.resolve) { + return undefined; + } + + if (!attachmentVersion) { + return undefined; + } + + const referenceData = attachmentVersion.raw_data ?? attachmentVersion.data; + const resolved = await definition.resolve( + { + id, + type, + data: referenceData, + }, + context + ); + + if (resolved === undefined) { + return undefined; + } + + const currentResolved = + attachmentVersion.raw_data !== undefined ? attachmentVersion.data : undefined; + const isSameResolved = + currentResolved !== undefined && JSON.stringify(currentResolved) === JSON.stringify(resolved); + + if (isSameResolved) { + return resolved; + } + + const newContentHash = hashContent(resolved); + const newTokens = estimateTokens(resolved); + + if (attachmentVersion.raw_data === undefined) { + // First resolve: cache resolved data on the existing version without bumping version. + attachmentVersion.raw_data = attachmentVersion.data; + attachmentVersion.data = resolved; + attachmentVersion.content_hash = newContentHash; + attachmentVersion.estimated_tokens = newTokens; + this.dirty = true; + return resolved; + } + + const newVersionNumber = attachment.current_version + 1; + const newVersion: AttachmentVersion = { + ...attachmentVersion, + version: newVersionNumber, + data: resolved, + raw_data: attachmentVersion.raw_data, + created_at: new Date().toISOString(), + content_hash: newContentHash, + estimated_tokens: newTokens, + }; + + attachment.versions.push(newVersion); + attachment.current_version = newVersionNumber; + this.dirty = true; + + return resolved; + } + getTotalTokenEstimate(): number { let total = 0; for (const attachment of this.attachments.values()) { diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/index.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/index.ts index c6ac7fddeab61..da58b43ee3c28 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/index.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/index.ts @@ -12,6 +12,7 @@ export type { AttachmentValidationResult, AgentFormattedAttachment, AttachmentFormatContext, + AttachmentResolveContext, } from './type_definition'; export type { AttachmentBoundedTool, diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/type_definition.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/type_definition.ts index 7ebc062ff0771..9bd7d6567e5ff 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/type_definition.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/attachments/type_definition.ts @@ -8,6 +8,7 @@ import type { MaybePromise } from '@kbn/utility-types'; import type { Attachment } from '@kbn/agent-builder-common/attachments'; import type { KibanaRequest } from '@kbn/core-http-server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import type { AttachmentBoundedTool } from './tools'; /** @@ -29,6 +30,19 @@ export interface AttachmentTypeDefinition, context: AttachmentFormatContext ) => MaybePromise; + /** + * Optional hook to resolve additional data on-demand when an attachment is read. + * + * This is primarily intended for "by-reference" attachment types, where the stored data is a + * reference to another entity (e.g. a saved object), and the actual content is resolved at read time. + * + * The return value is intentionally generic and will be exposed under a `resolved` key by + * `attachment_read` (and related server APIs). + */ + resolve?: ( + attachment: Attachment, + context: AttachmentResolveContext + ) => MaybePromise; /** * should return the list of tools from the registry which should be exposed to the agent * when attachments of that type are present in the conversation. @@ -58,6 +72,17 @@ export interface AttachmentFormatContext { spaceId: string; } +/** + * Context passed to the {@link AttachmentTypeDefinition.resolve} hook. + */ +export interface AttachmentResolveContext extends AttachmentFormatContext { + /** + * Saved objects client scoped to the current user. + * Optional to keep the core attachment contract generic and allow non-Kibana environments. + */ + savedObjectsClient?: SavedObjectsClientContract; +} + /** * Return type for attachment's validation handlers. * Refer to {@link InlineAttachmentTypeDefinition.validate} diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/moon.yml b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/moon.yml index bc4856941f6df..42c7eabe6996a 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/moon.yml +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/moon.yml @@ -27,6 +27,7 @@ dependsOn: - '@kbn/inference-common' - '@kbn/logging' - '@kbn/core-ui-settings-server' + - '@kbn/core-saved-objects-api-server' tags: - shared-common - package diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tools/handler.ts b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tools/handler.ts index 689393371919a..970f928d7cdd9 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tools/handler.ts +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tools/handler.ts @@ -9,6 +9,7 @@ import type { MaybePromise } from '@kbn/utility-types'; import type { Logger } from '@kbn/logging'; import type { IScopedClusterClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import type { ToolResult } from '@kbn/agent-builder-common/tools/tool_result'; import type { PromptRequest } from '@kbn/agent-builder-common/agents/prompts'; import type { @@ -88,6 +89,10 @@ export interface ToolHandlerContext { * Can be used to access ES on behalf of either the current user or the system user. */ esClient: IScopedClusterClient; + /** + * Saved objects client scoped to the current user. + */ + savedObjectsClient: SavedObjectsClientContract; /** * Inference model provider scoped to the current user. * Can be used to access the inference APIs or chatModel. diff --git a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tsconfig.json b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tsconfig.json index 96247b3eacb80..cc2effda84a86 100644 --- a/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tsconfig.json +++ b/x-pack/platform/packages/shared/agent-builder/agent-builder-server/tsconfig.json @@ -23,5 +23,6 @@ "@kbn/inference-common", "@kbn/logging", "@kbn/core-ui-settings-server", + "@kbn/core-saved-objects-api-server", ] } diff --git a/x-pack/platform/plugins/shared/agent_builder/server/routes/attachments.ts b/x-pack/platform/plugins/shared/agent_builder/server/routes/attachments.ts index 51afc8631f9d7..a8a896c2250ae 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/routes/attachments.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/routes/attachments.ts @@ -55,6 +55,20 @@ function isAttachmentReferencedInRounds( return false; } +const hasClientId = (attachment: { client_id?: string; versions: Array<{ data: unknown }> }) => { + if (attachment.client_id) { + return true; + } + + return attachment.versions.some((version) => { + if (!version?.data || typeof version.data !== 'object') { + return false; + } + + return Boolean((version.data as { client_id?: string }).client_id); + }); +}; + export function registerAttachmentRoutes({ router, getInternalServices, @@ -200,7 +214,7 @@ export function registerAttachmentRoutes({ }); // Check for duplicate ID if provided - if (id && stateManager.get(id)) { + if (id && stateManager.getAttachmentRecord(id)) { return response.conflict({ body: { message: `Attachment with ID '${id}' already exists` }, }); @@ -286,7 +300,7 @@ export function registerAttachmentRoutes({ const stateManager = createAttachmentStateManager(conversation.attachments ?? [], { getTypeDefinition: attachmentsService.getTypeDefinition, }); - const existing = stateManager.get(attachmentId); + const existing = stateManager.getAttachmentRecord(attachmentId); if (!existing) { return response.notFound({ @@ -393,7 +407,7 @@ export function registerAttachmentRoutes({ const stateManager = createAttachmentStateManager(conversation.attachments ?? [], { getTypeDefinition: attachmentsService.getTypeDefinition, }); - const existing = stateManager.get(attachmentId); + const existing = stateManager.getAttachmentRecord(attachmentId); if (!existing) { return response.notFound({ @@ -409,22 +423,19 @@ export function registerAttachmentRoutes({ } if (permanent) { - // Check if attachment is referenced in rounds - if (isAttachmentReferencedInRounds(attachmentId, conversation.rounds)) { + if (hasClientId(existing)) { return response.conflict({ body: { - message: `Cannot permanently delete attachment '${attachmentId}' because it is referenced in conversation rounds`, + message: `Cannot permanently delete attachment '${attachmentId}' because it was created from flyout configuration`, }, }); } - // Check if attachment has client_id (from flyout config) - const latestVersion = stateManager.getLatest(attachmentId); - const versionData = latestVersion?.data as Record | undefined; - if (versionData?.client_id) { + // Check if attachment is referenced in rounds + if (isAttachmentReferencedInRounds(attachmentId, conversation.rounds)) { return response.conflict({ body: { - message: `Cannot permanently delete attachment '${attachmentId}' because it was created from flyout configuration`, + message: `Cannot permanently delete attachment '${attachmentId}' because it is referenced in conversation rounds`, }, }); } @@ -516,7 +527,7 @@ export function registerAttachmentRoutes({ const stateManager = createAttachmentStateManager(conversation.attachments ?? [], { getTypeDefinition: attachmentsService.getTypeDefinition, }); - const existing = stateManager.get(attachmentId); + const existing = stateManager.getAttachmentRecord(attachmentId); if (!existing) { return response.notFound({ @@ -538,7 +549,7 @@ export function registerAttachmentRoutes({ }); } - const restored = stateManager.get(attachmentId)!; + const restored = stateManager.getAttachmentRecord(attachmentId)!; // Save the updated conversation await client.update({ @@ -609,7 +620,7 @@ export function registerAttachmentRoutes({ const stateManager = createAttachmentStateManager(conversation.attachments ?? [], { getTypeDefinition: attachmentsService.getTypeDefinition, }); - const existing = stateManager.get(attachmentId); + const existing = stateManager.getAttachmentRecord(attachmentId); if (!existing) { return response.notFound({ @@ -625,7 +636,7 @@ export function registerAttachmentRoutes({ }); } - const renamed = stateManager.get(attachmentId)!; + const renamed = stateManager.getAttachmentRecord(attachmentId)!; // Save the updated conversation await client.update({ diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/attachment_presentation.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/attachment_presentation.ts index 9f7cd49520641..367045a2fa178 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/attachment_presentation.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/attachment_presentation.ts @@ -180,7 +180,6 @@ const formatSummaryAttachments = (attachments: VersionedAttachment[]): string => /** * Formats attachment content based on type. - * Special handling for visualization_ref to show reference info instead of full resolved content. */ const formatAttachmentContent = (attachment: VersionedAttachment, data: unknown): string => { if (typeof data === 'string') { diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/prepare_conversation.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/prepare_conversation.ts index 5fb1f674bc6c8..4e9aae8d192fb 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/prepare_conversation.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/agents/modes/utils/prepare_conversation.ts @@ -88,7 +88,7 @@ const mergeInputAttachmentsIntoAttachmentState = async ( for (const input of inputs) { // Prefer stable IDs (if provided) if (input.id) { - const existing = attachmentStateManager.get(input.id); + const existing = attachmentStateManager.getAttachmentRecord(input.id); if (existing) { await attachmentStateManager.update( input.id, @@ -197,6 +197,16 @@ export const prepareConversation = async ({ }) ); + const activeAttachments = attachmentStateManager.getActive(); + await Promise.all( + activeAttachments.map(async (attachment) => { + await attachmentStateManager.get(attachment.id, { + version: attachment.current_version, + context, + }); + }) + ); + const versionedAttachmentPresentation = await prepareAttachmentPresentation( attachmentStateManager.getAll(), undefined, @@ -220,11 +230,11 @@ export const prepareConversation = async ({ }, formatContext ); - if (!formatted.getRepresentation) { + if (!formatted?.getRepresentation) { return undefined; } const representation = await formatted.getRepresentation(); - return representation.type === 'text' ? representation.value : undefined; + return representation?.type === 'text' ? representation.value : undefined; } catch { return undefined; } @@ -298,12 +308,21 @@ const prepareAttachment = async ({ try { const formatted = await definition.format(attachment, formatContext); const tools = formatted.getBoundedTools ? await formatted.getBoundedTools() : []; + if (!formatted.getRepresentation) { + return { + attachment, + representation: { type: 'text', value: JSON.stringify(attachment.data) }, + tools, + }; + } + const baseRepresentation = await formatted.getRepresentation(); + const representation = definition.resolve + ? withByReferenceNote({ representation: baseRepresentation, attachment }) + : baseRepresentation; return { attachment, - representation: formatted.getRepresentation - ? await formatted.getRepresentation() - : { type: 'text', value: JSON.stringify(attachment.data) }, + representation, tools, }; } catch (e) { @@ -315,6 +334,43 @@ const prepareAttachment = async ({ } }; +const withByReferenceNote = ({ + representation, + attachment, +}: { + representation: AttachmentRepresentation; + attachment: Attachment; +}): AttachmentRepresentation => { + if (representation.type !== 'text') { + return representation; + } + + const parts: string[] = []; + const note = + 'Note: this attachment is by-reference. Use attachment_read to resolve the full content.'; + + const trimmedValue = representation.value?.trim(); + if (trimmedValue) { + parts.push(trimmedValue); + } + + try { + parts.push(`Attachment data:\n${JSON.stringify(attachment.data, null, 2)}`); + } catch (e) { + // ignore stringify errors; fallback to whatever was already present + } + + // Avoid duplicating the note if the formatter already mentioned attachment_read. + if (!representation.value?.includes('attachment_read')) { + parts.push(note); + } + + return { + ...representation, + value: parts.filter(Boolean).join('\n\n'), + }; +}; + const inputToFinal = (input: AttachmentInput): Attachment => { return { ...input, diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/run_agent.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/run_agent.ts index 17a3ed37c5e09..355639f6d783a 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/run_agent.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/run_agent.ts @@ -33,6 +33,7 @@ export const createAgentHandlerContext = async const { request, elasticsearch, + savedObjects, spaces, modelProvider, toolsService, @@ -182,6 +183,7 @@ export const createToolHandlerContext = async spaceId, logger, esClient: elasticsearch.client.asScoped(request), + savedObjectsClient: savedObjects.getScopedClient(request), modelProvider, runner: manager.getRunner(), toolProvider: createToolProvider({ diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner.ts index a53350d659a61..2c2edef28911c 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner.ts @@ -9,6 +9,7 @@ import type { Logger } from '@kbn/logging'; import type { ElasticsearchServiceStart } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { SecurityServiceStart } from '@kbn/core-security-server'; +import type { SavedObjectsServiceStart } from '@kbn/core-saved-objects-server'; import type { SpacesPluginStart } from '@kbn/spaces-plugin/server'; import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server'; import { isAgentBuilderError, createInternalError } from '@kbn/agent-builder-common'; @@ -48,6 +49,7 @@ export interface CreateScopedRunnerDeps { // core services elasticsearch: ElasticsearchServiceStart; security: SecurityServiceStart; + savedObjects: SavedObjectsServiceStart; // external plugin deps spaces: SpacesPluginStart | undefined; actions: ActionsPluginStart; diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner_factory.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner_factory.ts index 74eb0e9f4037d..fa2e6df402cd9 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner_factory.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/runner/runner_factory.ts @@ -25,6 +25,7 @@ export class RunnerFactoryImpl implements RunnerFactory { const { inference, trackingService, uiSettings, savedObjects, ...otherDeps } = this.deps; return { ...otherDeps, + savedObjects, trackingService, modelProviderFactory: createModelProviderFactory({ inference, diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_add.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_add.ts index abf18a7733ca2..ffbb90675618f 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_add.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_add.ts @@ -50,7 +50,8 @@ export const createAttachmentAddTool = ({ } // Check for duplicate ID if provided - if (id && attachmentManager.get(id)) { + const existing = id ? attachmentManager.getAttachmentRecord(id) : undefined; + if (existing) { return { results: [ { diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_diff.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_diff.ts index b84f8b38ec00f..4feed37175bd6 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_diff.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_diff.ts @@ -7,7 +7,7 @@ import { z } from '@kbn/zod'; import { platformCoreTools, ToolType } from '@kbn/agent-builder-common'; -import { ATTACHMENT_REF_ACTOR } from '@kbn/agent-builder-common/attachments'; +import { getVersion } from '@kbn/agent-builder-common/attachments'; import { ToolResultType, isOtherResult } from '@kbn/agent-builder-common/tools/tool_result'; import type { BuiltinToolDefinition } from '@kbn/agent-builder-server'; import { createErrorResult, getToolResultId } from '@kbn/agent-builder-server'; @@ -37,7 +37,7 @@ export const createAttachmentDiffTool = ({ from_version: fromVersion, to_version: toVersion, }) => { - const attachment = attachmentManager.get(attachmentId); + const attachment = attachmentManager.getAttachmentRecord(attachmentId); if (!attachment) { return { @@ -67,16 +67,8 @@ export const createAttachmentDiffTool = ({ }; } - const fromVersionData = attachmentManager.readVersion( - attachmentId, - fromVersion, - ATTACHMENT_REF_ACTOR.agent - ); - const toVersionData = attachmentManager.readVersion( - attachmentId, - toVersion, - ATTACHMENT_REF_ACTOR.agent - ); + const fromVersionData = getVersion(attachment, fromVersion); + const toVersionData = getVersion(attachment, toVersion); return { results: [ diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_read.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_read.ts index 1df57a92267c2..008c818918b68 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_read.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_read.ts @@ -7,7 +7,6 @@ import { z } from '@kbn/zod'; import { platformCoreTools, ToolType } from '@kbn/agent-builder-common'; -import { ATTACHMENT_REF_ACTOR } from '@kbn/agent-builder-common/attachments'; import { ToolResultType, isOtherResult } from '@kbn/agent-builder-common/tools/tool_result'; import type { BuiltinToolDefinition } from '@kbn/agent-builder-server'; import { createErrorResult, getToolResultId } from '@kbn/agent-builder-server'; @@ -37,8 +36,11 @@ export const createAttachmentReadTool = ({ 'Read the content of a conversation attachment by ID. Use this to retrieve data you previously stored or to check the current state of an attachment.', schema: attachmentReadSchema, tags: ['attachment'], - handler: async ({ attachment_id: attachmentId, version }) => { - const attachment = attachmentManager.get(attachmentId); + handler: async ({ attachment_id: attachmentId, version }, context) => { + const attachment = await attachmentManager.get(attachmentId, { + version, + context, + }); if (!attachment) { return { @@ -51,27 +53,14 @@ export const createAttachmentReadTool = ({ }; } - const versionData = version - ? attachmentManager.readVersion(attachmentId, version, ATTACHMENT_REF_ACTOR.agent) - : attachmentManager.readLatest(attachmentId, ATTACHMENT_REF_ACTOR.agent); + const { data: versionData, type } = attachment; - if (!versionData) { - return { - results: [ - createErrorResult({ - message: `Version ${version} not found for attachment '${attachmentId}'`, - metadata: { attachment_id: attachmentId, version }, - }), - ], - }; - } - - let data = versionData.data; + let formattedData: unknown = versionData.data; + const rawData = (versionData as { raw_data?: unknown }).raw_data; if (attachmentsService && formatContext) { const definition = attachmentsService.getTypeDefinition(attachment.type); const typeReadonly = definition?.isReadonly ?? true; - const isReadonly = typeReadonly || attachment.readonly === true; - if (definition && isReadonly) { + if (definition && typeReadonly) { try { const formatted = await definition.format( { @@ -83,13 +72,13 @@ export const createAttachmentReadTool = ({ ); if (formatted.getRepresentation) { const representation = await formatted.getRepresentation(); - data = + formattedData = representation.type === 'text' ? representation.value : JSON.stringify(representation); } } catch { - data = versionData.data; + formattedData = versionData.data; } } } @@ -101,9 +90,10 @@ export const createAttachmentReadTool = ({ type: ToolResultType.other, data: { attachment_id: attachmentId, - type: attachment.type, - version: versionData.version, - data, + type, + version: attachment.version, + data: formattedData, + ...(rawData !== undefined ? { raw_data: rawData } : {}), }, }, ], diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_tools.test.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_tools.test.ts index 95606d613c3f5..7922b0ebd9880 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_tools.test.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_tools.test.ts @@ -6,10 +6,16 @@ */ import { ToolResultType } from '@kbn/agent-builder-common'; -import { httpServerMock } from '@kbn/core-http-server-mocks'; +import type { Attachment } from '@kbn/agent-builder-common/attachments'; +import { AttachmentType } from '@kbn/agent-builder-common/attachments'; +import type { + AttachmentResolveContext, + AttachmentTypeDefinition, +} from '@kbn/agent-builder-server/attachments'; import { createAttachmentStateManager } from '@kbn/agent-builder-server/attachments'; import type { AttachmentStateManager } from '@kbn/agent-builder-server/attachments'; import type { ToolHandlerStandardReturn } from '@kbn/agent-builder-server/tools'; +import { httpServerMock } from '@kbn/core-http-server-mocks'; import { createAttachmentTools } from '.'; describe('attachment tools', () => { @@ -153,39 +159,68 @@ describe('attachment tools', () => { expect((result.results[0] as any).data.data).toBe('hello'); }); - it('returns formatted content for read-only attachments', async () => { - const readonlyAttachmentsService = { - getTypeDefinition: () => ({ - id: 'text', - validate: (input: unknown) => ({ valid: true, data: input }), - format: (attachment: { data: unknown }) => ({ - getRepresentation: () => ({ - type: 'text', - value: `formatted:${String(attachment.data)}`, + it('resolves visualization_ref attachments when savedObjectsClient is available', async () => { + const customAttachmentsService = { + getTypeDefinition: () => + ({ + id: AttachmentType.visualizationRef, + validate: (input: unknown) => ({ valid: true, data: input }), + format: (formattedAttachment: Attachment) => ({ + getRepresentation: () => ({ + type: 'text', + value: JSON.stringify(formattedAttachment.data), + }), }), - }), - isReadonly: true, - }), + resolve: async (_a: Attachment, ctx: AttachmentResolveContext) => { + const resolved = await ctx.savedObjectsClient!.resolve('lens', 'so-123'); + return { + found: true, + outcome: resolved.outcome, + alias_target_id: resolved.alias_target_id, + saved_object_id: resolved.saved_object.id, + saved_object_type: resolved.saved_object.type, + updated_at: resolved.saved_object.updated_at, + attributes: resolved.saved_object.attributes, + title: (resolved.saved_object.attributes as any).title, + description: (resolved.saved_object.attributes as any).description, + }; + }, + } as unknown as AttachmentTypeDefinition), } as any; - + const resolveAttachmentManager = createAttachmentStateManager([], { + getTypeDefinition: customAttachmentsService.getTypeDefinition, + }); + const attachment = await resolveAttachmentManager.add({ + type: AttachmentType.visualizationRef, + data: { + saved_object_id: 'so-123', + }, + description: 'Lens ref', + }); const tool = createAttachmentTools({ - attachmentManager, - attachmentsService: readonlyAttachmentsService, + attachmentManager: resolveAttachmentManager, + attachmentsService: customAttachmentsService, formatContext, }).find((t) => t.id === 'platform.core.attachment_read')!; + const result = (await tool.handler({ attachment_id: attachment.id }, { + savedObjectsClient: { + resolve: async () => ({ + outcome: 'exactMatch', + alias_target_id: null, + saved_object: { + id: 'so-123', + type: 'lens', + updated_at: '2026-01-01T00:00:00.000Z', + attributes: { title: 'My Lens', description: 'Desc', state: { a: 1 } }, + }, + }), + }, + } as any)) as ToolHandlerStandardReturn; - const attachment = await attachmentManager.add({ - type: 'text', - data: 'hello', - description: 'Test', - }); - - const result = (await tool.handler( - { attachment_id: attachment.id }, - {} as any - )) as ToolHandlerStandardReturn; - - expect((result.results[0] as any).data.data).toBe('formatted:hello'); + expect((result.results[0] as any).data.type).toBe(AttachmentType.visualizationRef); + expect((result.results[0] as any).data.raw_data).toEqual({ saved_object_id: 'so-123' }); + expect((result.results[0] as any).data.data).toContain('"found":true'); + expect((result.results[0] as any).data.data).toContain('"title":"My Lens"'); }); it('reads a specific version', async () => { diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_update.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_update.ts index dd581f5686efb..11320533fba66 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_update.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/services/tools/builtin/attachments/attachment_update.ts @@ -34,7 +34,7 @@ export const createAttachmentUpdateTool = ({ schema: attachmentUpdateSchema, tags: ['attachment'], handler: async ({ attachment_id: attachmentId, data, description }) => { - const existing = attachmentManager.get(attachmentId); + const existing = attachmentManager.getAttachmentRecord(attachmentId); if (!existing) { return { diff --git a/x-pack/platform/plugins/shared/agent_builder/server/test_utils/runner.ts b/x-pack/platform/plugins/shared/agent_builder/server/test_utils/runner.ts index 4d088d7dc52f6..528e47156e600 100644 --- a/x-pack/platform/plugins/shared/agent_builder/server/test_utils/runner.ts +++ b/x-pack/platform/plugins/shared/agent_builder/server/test_utils/runner.ts @@ -10,6 +10,7 @@ import { loggerMock } from '@kbn/logging-mocks'; import { elasticsearchServiceMock, httpServerMock, + savedObjectsServiceMock, securityServiceMock, } from '@kbn/core/server/mocks'; import { spacesMock } from '@kbn/spaces-plugin/server/mocks'; @@ -103,10 +104,7 @@ export const createStateManagerMock = (): StateManagerMock => { export const createAttachmentStateManagerMock = (): AttachmentStateManagerMock => { return { get: jest.fn(), - getLatest: jest.fn(), - getVersion: jest.fn(), - readLatest: jest.fn(), - readVersion: jest.fn(), + getAttachmentRecord: jest.fn(), getActive: jest.fn(), getAll: jest.fn(), getDiff: jest.fn(), @@ -119,6 +117,7 @@ export const createAttachmentStateManagerMock = (): AttachmentStateManagerMock = getAccessedRefs: jest.fn(), clearAccessTracking: jest.fn(), resolveRefs: jest.fn(), + resolveAttachment: jest.fn(), getTotalTokenEstimate: jest.fn(), hasChanges: jest.fn(), markClean: jest.fn(), @@ -180,6 +179,7 @@ export const createAgentHandlerContextMock = (): AgentHandlerContextMock => { request: httpServerMock.createKibanaRequest(), spaceId: 'default', esClient: elasticsearchServiceMock.createScopedClusterClient(), + savedObjectsClient: savedObjectsServiceMock.createStartContract().getScopedClient({} as any), modelProvider: createModelProviderMock(), toolProvider: createToolProviderMock(), runner: createScopedRunnerMock(), @@ -203,6 +203,9 @@ export interface ToolHandlerContextMock extends ToolHandlerContext { filestore: FileSystemStoreMock; prompts: ToolPromptManagerMock; stateManager: ToolStateManagerMock; + savedObjectsClient: ReturnType< + ReturnType['getScopedClient'] + >; } export const createToolHandlerContextMock = (): ToolHandlerContextMock => { @@ -223,6 +226,7 @@ export const createToolHandlerContextMock = (): ToolHandlerContextMock => { stateManager: createToolStateManagerMock(), attachments: createAttachmentStateManagerMock(), filestore: createFileSystemStoreMock(), + savedObjectsClient: savedObjectsServiceMock.createStartContract().getScopedClient({} as any), }; }; @@ -238,6 +242,7 @@ export const createScopedRunnerDepsMock = (): CreateScopedRunnerDepsMock => { return { elasticsearch: elasticsearchServiceMock.createStart(), security: securityServiceMock.createStart(), + savedObjects: savedObjectsServiceMock.createStartContract(), spaces: spacesMock.createStart(), actions: actionsMock.createStart(), modelProvider: createModelProviderMock(), @@ -258,6 +263,7 @@ export const createRunnerDepsMock = (): CreateRunnerDepsMock => { return { elasticsearch: elasticsearchServiceMock.createStart(), security: securityServiceMock.createStart(), + savedObjects: savedObjectsServiceMock.createStartContract(), spaces: spacesMock.createStart(), actions: actionsMock.createStart(), modelProviderFactory: createModelProviderFactoryMock(), diff --git a/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/index.ts b/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/index.ts index 45fce1fcee61f..8edecb86a7d09 100644 --- a/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/index.ts +++ b/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/index.ts @@ -10,6 +10,7 @@ import type { CoreSetup } from '@kbn/core-lifecycle-server'; import { createTextAttachmentType } from './text'; import { createEsqlAttachmentType } from './esql'; import { createScreenContextAttachmentType } from './screen_context'; +import { createVisualizationRefAttachmentType } from './visualization_ref'; import type { AgentBuilderPlatformPluginStart, PluginSetupDependencies, @@ -29,6 +30,7 @@ export const registerAttachmentTypes = ({ createTextAttachmentType(), createScreenContextAttachmentType(), createEsqlAttachmentType(), + createVisualizationRefAttachmentType(), ]; attachmentTypes.forEach((attachmentType) => { diff --git a/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/visualization_ref.ts b/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/visualization_ref.ts new file mode 100644 index 0000000000000..17b7e24518783 --- /dev/null +++ b/x-pack/platform/plugins/shared/agent_builder_platform/server/attachment_types/visualization_ref.ts @@ -0,0 +1,93 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { VisualizationRefAttachmentData } from '@kbn/agent-builder-common/attachments'; +import { + AttachmentType, + visualizationRefAttachmentDataSchema, +} from '@kbn/agent-builder-common/attachments'; +import type { + AttachmentResolveContext, + AttachmentTypeDefinition, +} from '@kbn/agent-builder-server/attachments'; +import { + LensConfigBuilder, + type LensApiSchemaType, + type LensAttributes, +} from '@kbn/lens-embeddable-utils/config_builder'; + +/** + * Creates the definition for the `visualization_ref` attachment type. + * + * This attachment type represents a reference to a saved Lens visualization object. + * It stores the saved object ID and type, + * with optional cached metadata (title, description). + * + * The referenced content is resolved on-demand on the server (e.g. via attachment_read tool). + */ +export const createVisualizationRefAttachmentType = (): AttachmentTypeDefinition< + AttachmentType.visualizationRef, + VisualizationRefAttachmentData +> => { + return { + id: AttachmentType.visualizationRef, + validate: (input) => { + const parseResult = visualizationRefAttachmentDataSchema.safeParse(input); + if (parseResult.success) { + return { valid: true, data: parseResult.data }; + } else { + return { valid: false, error: parseResult.error.message }; + } + }, + format: (attachment) => ({ + getRepresentation: () => ({ + // Keep formatting minimal; generic by-ref messaging is added centrally. + type: 'text', + value: JSON.stringify(attachment.data, null, 2), + }), + }), + resolve: async (attachment, context: AttachmentResolveContext) => { + // Allow this type to be used in environments where SO client isn't available. + if (!context.savedObjectsClient) return undefined; + + const { saved_object_id } = attachment.data; + + try { + const resolveResult = await context.savedObjectsClient.resolve('lens', saved_object_id); + const savedObject = resolveResult.saved_object as { error?: { message?: string } }; + + if (savedObject?.error) { + return undefined; + } + + const lensAttributes = toLensAttributes( + resolveResult.saved_object.attributes as LensAttributes, + resolveResult.saved_object.references + ); + + return toLensApiConfig(lensAttributes); + } catch (error) { + return undefined; + } + }, + getAgentDescription: () => { + return `A visualization_ref attachment contains a reference to a saved Lens visualization. The reference includes the saved object ID and type. Use attachment_read to resolve the referenced content when needed.`; + }, + getTools: () => [], + }; +}; + +const toLensAttributes = ( + attributes: LensAttributes, + references: LensAttributes['references'] | undefined +): LensAttributes => ({ + ...attributes, + references: references ?? attributes.references ?? [], +}); + +const toLensApiConfig = (attributes: LensAttributes): LensApiSchemaType => + new LensConfigBuilder().toAPIFormat(attributes); diff --git a/x-pack/platform/plugins/shared/agent_builder_platform/server/tools/create_visualization/create_visualization.ts b/x-pack/platform/plugins/shared/agent_builder_platform/server/tools/create_visualization/create_visualization.ts index 0d5b5c50c8304..710b71bb88d45 100644 --- a/x-pack/platform/plugins/shared/agent_builder_platform/server/tools/create_visualization/create_visualization.ts +++ b/x-pack/platform/plugins/shared/agent_builder_platform/server/tools/create_visualization/create_visualization.ts @@ -18,6 +18,7 @@ import { xyStateSchema } from '@kbn/lens-embeddable-utils/config_builder/schema/ import { regionMapStateSchemaESQL } from '@kbn/lens-embeddable-utils/config_builder/schema/charts/region_map'; import { heatmapStateSchemaESQL } from '@kbn/lens-embeddable-utils/config_builder/schema/charts/heatmap'; import { getToolResultId } from '@kbn/agent-builder-server'; +import { getLatestVersion } from '@kbn/agent-builder-common/attachments'; import { AGENT_BUILDER_DASHBOARD_TOOLS_SETTING_ID } from '@kbn/management-settings-ids'; import type { VisualizationConfig } from './types'; import { guessChartType } from './guess_chart_type'; @@ -103,9 +104,9 @@ This tool will: let parsedExistingConfig: VisualizationConfig | null = null; if (attachmentId) { - const existingAttachment = attachments.get(attachmentId); - if (existingAttachment) { - const latestVersion = attachments.getLatest(attachmentId); + const existingAttachmentRecord = attachments.getAttachmentRecord(attachmentId); + if (existingAttachmentRecord) { + const latestVersion = getLatestVersion(existingAttachmentRecord); if (latestVersion?.data) { parsedExistingConfig = latestVersion.data as VisualizationConfig; existingConfig = JSON.stringify(parsedExistingConfig); @@ -184,7 +185,7 @@ This tool will: let version: number; let isUpdate = false; - if (attachmentId && attachments.get(attachmentId)) { + if (attachmentId && attachments.getAttachmentRecord(attachmentId)) { const updated = await attachments.update(attachmentId, { data: visualizationData, description: `Visualization: ${nlQuery.slice(0, 50)}${