diff --git a/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.spec.ts b/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.spec.ts new file mode 100644 index 00000000000..a67bd4cdb97 --- /dev/null +++ b/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.spec.ts @@ -0,0 +1,869 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { AssistantSendMessage, AssistantSocketEvent } from '@zeppelin/sdk'; + +import { + AssistantHttpError, + createAssistantTransport, + mapSocketEvent, + createRunTracking, + scopeTransport, + socketRun, + toVisibleMessages +} from './assistantTransport'; +import { + AssistantConnectionError, + AssistantStreamError, + type AssistantRunEvent, + type AssistantRunState +} from './assistantContract'; +import { joinPages } from './messageHistory'; + +/** A host socket double: records sends and lets a test emit server events. */ +const fakeSocket = () => { + const listeners = new Set<(event: AssistantSocketEvent) => void>(); + const closeListeners = new Set<() => void>(); + const sent: AssistantSendMessage[] = []; + return { + sent, + listenerCount: () => listeners.size, + emit: (event: AssistantSocketEvent) => listeners.forEach(listener => listener(event)), + send: vi.fn((message: AssistantSendMessage) => { + sent.push(message); + }), + subscribe: vi.fn((listener: (event: AssistantSocketEvent) => void) => { + listeners.add(listener); + return () => listeners.delete(listener); + }), + close: () => closeListeners.forEach(listener => listener()), + closeListenerCount: () => closeListeners.size, + subscribeClose: vi.fn((listener: () => void) => { + closeListeners.add(listener); + return () => closeListeners.delete(listener); + }) + }; +}; + +const forNote = (noteId: string, onAuthError?: (status: number, location: string | null) => void) => + createAssistantTransport('https://example.test/zeppelin/api', noteId, fakeSocket(), onAuthError); + +const iterate = (iterable: AsyncIterable): AsyncIterator => iterable[Symbol.asyncIterator](); + +const collect = async (events: AsyncIterable): Promise => { + const collected: AssistantRunEvent[] = []; + for await (const event of events) { + collected.push(event); + } + return collected; +}; + +const message = (content = 'hi'): AssistantSendMessage => ({ noteId: 'n', conversationId: 'c1', content }); + +describe('assistant transport', () => { + // One note's sends share their run tracking, as the transport does. + let tracking = createRunTracking(); + beforeEach(() => { + tracking = createRunTracking(); + }); + afterEach(() => { + vi.unstubAllGlobals(); + vi.useRealTimers(); + }); + + it('lists the most recently active conversation first and undated ones last', async () => { + const body = [ + { id: 'old', createdAt: '2026-10-01T00:00:00Z', updatedAt: '2026-10-01T00:00:00Z' }, + { id: 'undated' }, + { id: 'new', createdAt: '2026-09-30T00:00:00Z', updatedAt: '2026-10-03T00:00:00Z' }, + { id: 'created-only', createdAt: '2026-10-02T00:00:00Z' } + ]; + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue( + new Response(JSON.stringify({ status: 'OK', body }), { + status: 200, + headers: { 'Content-Type': 'application/json' } + }) + ) + ); + const conversations = await forNote('note-1').listConversations(); + expect(conversations.map(conversation => conversation.id)).toEqual(['new', 'created-only', 'old', 'undated']); + // The list shows the latest activity, or the creation time before any. + expect(conversations.map(conversation => conversation.updatedAt)).toEqual([ + '2026-10-03T00:00:00Z', + '2026-10-02T00:00:00Z', + '2026-10-01T00:00:00Z', + undefined + ]); + }); + + it('uses the per-note conversation endpoints and unwraps Zeppelin JSON responses', async () => { + const json = (body: unknown) => + new Response(JSON.stringify(body), { status: 200, headers: { 'Content-Type': 'application/json' } }); + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + json({ + status: 'OK', + body: [{ id: 'conv-1', title: 'First', noteId: 'note/one', ownerId: 'alice', canSendMessage: false }] + }) + ) + .mockResolvedValueOnce(json({ status: 'OK', body: { id: 'conv-2', title: 'Load the data', ownerId: 'bob' } })) + .mockResolvedValueOnce( + json({ + status: 'OK', + body: { + messages: [ + { id: 'message-2', role: 'assistant', content: 'hi there' }, + { id: 'message-1', role: 'user', content: 'hello' } + ], + cursor: null + } + }) + ) + .mockResolvedValueOnce(new Response(null, { status: 204 })); + vi.stubGlobal('fetch', fetchMock); + const transport = forNote('note/one'); + + await expect(transport.listConversations()).resolves.toEqual([ + { id: 'conv-1', title: 'First', ownerId: 'alice', canSendMessage: false } + ]); + await expect(transport.createConversation({ title: ' Load the data ' })).resolves.toEqual({ + id: 'conv-2', + title: 'Load the data', + ownerId: 'bob' + }); + // The server pages latest first; the panel wants oldest first. + await expect(transport.getMessages('conv/1')).resolves.toEqual({ + messages: [ + { id: 'message-1', role: 'user', content: 'hello' }, + { id: 'message-2', role: 'assistant', content: 'hi there' } + ], + earlierCursor: null + }); + await expect(transport.deleteConversation('conv/1')).resolves.toBeUndefined(); + + const base = 'https://example.test/zeppelin/api/notes/note%2Fone/conversations'; + expect(fetchMock.mock.calls.map(([url, init]) => [url, init.method])).toEqual([ + [base, 'GET'], + [base, 'POST'], + [`${base}/conv%2F1/messages?limit=50`, 'GET'], + [`${base}/conv%2F1`, 'DELETE'] + ]); + expect(fetchMock.mock.calls[1][1]).toMatchObject({ + credentials: 'include', + headers: expect.objectContaining({ 'X-Requested-With': 'XMLHttpRequest', 'Content-Type': 'application/json' }), + // The first question names the conversation, with whitespace collapsed. + body: JSON.stringify({ title: 'Load the data' }) + }); + }); + + it('words a 404 on the conversation list as a server without the assistant', async () => { + vi.stubGlobal( + 'fetch', + vi.fn(async () => new Response(null, { status: 404 })) + ); + const transport = forNote('n'); + + await expect(transport.listConversations()).rejects.toThrow('The assistant is not available on this server'); + await expect(transport.createConversation({})).rejects.toThrow('The assistant is not available on this server'); + await expect(transport.getMessages('c')).rejects.toThrow('This conversation or note no longer exists.'); + }); + + it('hands 401 and 405 to the host, but not other failures', async () => { + const onAuthError = vi.fn(); + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockResolvedValueOnce(new Response(null, { status: 401, headers: { Location: '/login' } })) + .mockResolvedValueOnce(new Response(null, { status: 405 })) + .mockResolvedValueOnce(new Response(null, { status: 500 })) + ); + const transport = forNote('n', onAuthError); + + await expect(transport.listConversations()).rejects.toMatchObject({ status: 401, location: '/login' }); + await expect(transport.createConversation({})).rejects.toMatchObject({ status: 405 }); + await expect(transport.getMessages('c')).rejects.toBeInstanceOf(AssistantHttpError); + + expect(onAuthError.mock.calls).toEqual([ + [401, '/login'], + [405, null] + ]); + }); + + it('sends over the socket and yields only this conversation until a terminal event', async () => { + const socket = fakeSocket(); + const transport = createAssistantTransport('https://example.test/api', 'n', socket); + const run = iterate(transport.openRun('c1', { prompt: 'improve' }, new AbortController().signal)); + const first = run.next(); + await vi.waitFor(() => expect(socket.sent).toHaveLength(1)); + expect(socket.sent[0]).toEqual({ noteId: 'n', conversationId: 'c1', content: 'improve' }); + + socket.emit({ conversationId: 'other', type: 'message.delta', payload: { messageId: 'x', delta: 'no' } }); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'run_1' } }); + socket.emit({ + conversationId: 'c1', + type: 'tool_call.started', + payload: { toolCallId: 't1', name: 'list_paragraphs' } + }); + socket.emit({ + conversationId: 'c1', + type: 'tool_call.done', + payload: { toolCallId: 't1', result: { value: '[]' } } + }); + // A newer server may send types this client does not know. + socket.emit({ conversationId: 'c1', type: 'usage.extra', payload: {} } as unknown as AssistantSocketEvent); + socket.emit({ conversationId: 'c1', type: 'message.delta', payload: { messageId: 'm1', delta: '한글' } }); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'run_1' } }); + socket.emit({ conversationId: 'c1', type: 'message.delta', payload: { messageId: 'm2', delta: 'late' } }); + + const events = [(await first).value]; + for (let next = await run.next(); !next.done; next = await run.next()) events.push(next.value); + expect(events).toEqual([ + { type: 'run.started', runId: 'run_1' }, + { type: 'tool_call.started', toolCallId: 't1', name: 'list_paragraphs' }, + { type: 'tool_call.done', toolCallId: 't1', name: undefined }, + { type: 'message.delta', messageId: 'm1', delta: '한글' }, + { type: 'run.completed', runId: 'run_1' } + ]); + expect(socket.listenerCount()).toBe(0); + }); + + it('ends on a run failure that arrives before run.started', async () => { + const socket = fakeSocket(); + const pending = collect(socketRun(socket, message(), new AbortController().signal, tracking)); + socket.emit({ + conversationId: 'c1', + type: 'run.failed', + payload: { runId: 'run_x', error: { status: 403 } } + }); + await expect(pending).resolves.toEqual([ + { + type: 'run.failed', + runId: 'run_x', + message: 'You do not have permission to send messages to this conversation.' + } + ]); + }); + + it('blocks resending an aborted request until its delayed start is tracked', async () => { + const socket = fakeSocket(); + const leave = new AbortController(); + const first = iterate(socketRun(socket, message('first'), leave.signal, tracking)); + const response = first.next(); + leave.abort(); + await expect(response).rejects.toMatchObject({ name: 'AbortError' }); + await expect( + collect(socketRun(socket, message('second'), new AbortController().signal, tracking)) + ).resolves.toEqual([ + { type: 'run.failed', message: 'The previous request is still being confirmed. Try again shortly.' } + ]); + expect(socket.sent).toHaveLength(1); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'old' } }); + expect(socket.listenerCount()).toBe(1); + expect(socket.closeListenerCount()).toBe(1); + const second = collect(socketRun(socket, message('second'), new AbortController().signal, tracking)); + socket.emit({ + conversationId: 'c1', + type: 'message.delta', + payload: { messageId: 'old-answer', delta: 'old answer' } + }); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'old' } }); + socket.emit({ conversationId: 'c1', type: 'run.failed', payload: { runId: 'new', error: { status: 409 } } }); + await expect(second).resolves.toEqual([ + { type: 'run.failed', runId: 'new', message: 'This conversation is still answering. Try again when it finishes.' } + ]); + }); + + it('tracks a queued start even when the consumer aborts before reading it', async () => { + const socket = fakeSocket(); + const controller = new AbortController(); + const first = iterate(socketRun(socket, message(), controller.signal, tracking)); + const response = first.next(); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'queued' } }); + controller.abort(); + await expect(response).rejects.toMatchObject({ name: 'AbortError' }); + expect(tracking.abandoned.has('queued')).toBe(true); + expect(socket.listenerCount()).toBe(1); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'queued' } }); + expect(socket.listenerCount()).toBe(0); + }); + + it('shares pending requests across transports on the same host connection', async () => { + const connectionKey = {}; + const socket = { ...fakeSocket(), connectionKey }; + const old = createAssistantTransport('/api', 'note', socket); + const leave = new AbortController(); + const response = iterate(old.openRun('c1', { prompt: 'first' }, leave.signal)).next(); + leave.abort(); + await expect(response).rejects.toMatchObject({ name: 'AbortError' }); + const next = createAssistantTransport('/api', 'note', { ...socket }); + await expect(collect(next.openRun('c1', { prompt: 'second' }, new AbortController().signal))).resolves.toEqual([ + { type: 'run.failed', message: 'The previous request is still being confirmed. Try again shortly.' } + ]); + expect(socket.sent).toHaveLength(1); + socket.emit({ conversationId: 'c1', type: 'run.failed', payload: { runId: 'old', error: { status: 403 } } }); + expect(socket.listenerCount()).toBe(0); + const sent = collect(next.openRun('c1', { prompt: 'second' }, new AbortController().signal)); + socket.emit({ conversationId: 'c1', type: 'run.failed', payload: { runId: 'new', error: { status: 409 } } }); + await expect(sent).resolves.toHaveLength(1); + expect(socket.sent).toHaveLength(2); + }); + + it('releases unconfirmed sends when the connection closes', async () => { + const socket = fakeSocket(); + const controller = new AbortController(); + const response = iterate(socketRun(socket, message(), controller.signal, tracking)).next(); + controller.abort(); + await expect(response).rejects.toMatchObject({ name: 'AbortError' }); + socket.close(); + expect(socket.listenerCount()).toBe(0); + expect(socket.closeListenerCount()).toBe(0); + const reconnect = collect(socketRun(socket, message(), new AbortController().signal, tracking)); + socket.emit({ conversationId: 'c1', type: 'run.failed', payload: { runId: 'new', error: { status: 403 } } }); + await expect(reconnect).resolves.toHaveLength(1); + }); + + it('ignores the rest of an abandoned run when the next message is sent', async () => { + const socket = fakeSocket(); + const leave = new AbortController(); + const abandoned = iterate(socketRun(socket, message('first'), leave.signal, tracking)); + const firstEvent = abandoned.next(); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'run_old' } }); + expect((await firstEvent).value).toEqual({ type: 'run.started', runId: 'run_old' }); + leave.abort(); + await expect(abandoned.next()).rejects.toThrow(); + + // Sent while the abandoned run still streams: the server rejects it, and the old run's tail is not this run's. + const rejected = collect(socketRun(socket, message('second'), new AbortController().signal, tracking)); + socket.emit({ conversationId: 'c1', type: 'message.delta', payload: { messageId: 'old', delta: 'tail' } }); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'run_old' } }); + socket.emit({ + conversationId: 'c1', + type: 'run.failed', + payload: { runId: 'run_new', error: { status: 409 } } + }); + await expect(rejected).resolves.toEqual([ + { + type: 'run.failed', + runId: 'run_new', + message: 'This conversation is still answering. Try again when it finishes.' + } + ]); + + // Once the old run ended, the next message runs normally. + const next = collect(socketRun(socket, message('third'), new AbortController().signal, tracking)); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'run_3' } }); + socket.emit({ conversationId: 'c1', type: 'message.delta', payload: { messageId: 'm3', delta: 'ok' } }); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'run_3' } }); + await expect(next).resolves.toEqual([ + { type: 'run.started', runId: 'run_3' }, + { type: 'message.delta', messageId: 'm3', delta: 'ok' }, + { type: 'run.completed', runId: 'run_3' } + ]); + }); + + it('words run.failed by its HTTP status and falls back for the rest', () => { + const failed = (error: unknown) => + mapSocketEvent({ conversationId: 'c', type: 'run.failed', payload: { runId: 'r', error } }); + expect(failed({ status: 409 })).toEqual({ + type: 'run.failed', + runId: 'r', + message: 'This conversation is still answering. Try again when it finishes.' + }); + expect(failed({ status: 400 })).toMatchObject({ + message: 'The assistant could not accept this message. Edit it and try again.' + }); + expect(failed({ status: 503 })).toMatchObject({ message: expect.stringContaining('not configured') }); + expect(failed({ status: 500 })).toMatchObject({ + message: 'The assistant hit an error while answering. Try again.' + }); + expect(failed(undefined)).toMatchObject({ message: 'The assistant hit an error while answering. Try again.' }); + expect(failed(null)).toMatchObject({ message: 'The assistant hit an error while answering. Try again.' }); + expect(failed(403)).toMatchObject({ message: 'The assistant hit an error while answering. Try again.' }); + expect(failed({ status: '403' })).toMatchObject({ + message: 'The assistant hit an error while answering. Try again.' + }); + }); + + it('loads an earlier page with the cursor and joins a turn split across pages', async () => { + const fetchMock = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + status: 'OK', + body: { messages: [{ id: 'm3', role: 'assistant', content: 'part one.' }], cursor: 'm3' } + }), + { status: 200, headers: { 'Content-Type': 'application/json' } } + ) + ); + vi.stubGlobal('fetch', fetchMock); + const page = await forNote('n').getMessages('c', 'm9'); + expect(fetchMock.mock.calls[0][0]).toBe( + 'https://example.test/zeppelin/api/notes/n/conversations/c/messages?limit=50&cursor=m9' + ); + expect(page).toEqual({ messages: [{ id: 'm3', role: 'assistant', content: 'part one.' }], earlierCursor: 'm3' }); + expect( + joinPages(page.messages, [ + { id: 'm5', role: 'assistant', content: 'part two.' }, + { id: 'm6', role: 'user', content: 'next' } + ]) + ).toEqual([ + { id: 'm5', role: 'assistant', content: 'part one.\n\npart two.' }, + { id: 'm6', role: 'user', content: 'next' } + ]); + }); + + it('explains a missing assistant configuration instead of the bare status', () => { + expect(new AssistantHttpError(503, null).message).toBe( + 'The assistant is not configured on this server. Ask an administrator to enable it.' + ); + expect(new AssistantHttpError(500, null).message).toBe('Assistant request failed with HTTP 500'); + }); + + it('keeps only the ACK listener when aborted before the run starts', async () => { + const socket = fakeSocket(); + const controller = new AbortController(); + const run = iterate(socketRun(socket, message(), controller.signal, tracking)); + const pending = run.next(); + await vi.waitFor(() => expect(socket.listenerCount()).toBe(2)); + + controller.abort(); + + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }); + expect(socket.listenerCount()).toBe(1); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'late' } }); + expect(socket.listenerCount()).toBe(1); + socket.emit({ conversationId: 'c1', type: 'run.completed', payload: { runId: 'late' } }); + expect(socket.listenerCount()).toBe(0); + }); + + it('ends a run with a connection error when the notebook socket closes', async () => { + const socket = fakeSocket(); + const pending = collect(socketRun(socket, message(), new AbortController().signal, tracking)); + socket.emit({ conversationId: 'c1', type: 'run.started', payload: { runId: 'r1' } }); + socket.close(); + await expect(pending).rejects.toBeInstanceOf(AssistantConnectionError); + expect(socket.closeListenerCount()).toBe(0); + }); + + it('reports a retryable error when the run goes quiet past the idle timeout', async () => { + vi.useFakeTimers(); + const socket = fakeSocket(); + const pending = collect(socketRun(socket, message(), new AbortController().signal, tracking, 1000)); + const outcome = pending.then( + () => null, + (error: unknown) => error + ); + await vi.advanceTimersByTimeAsync(1000); + expect(await outcome).toBeInstanceOf(AssistantStreamError); + expect(socket.listenerCount()).toBe(1); + socket.close(); + expect(socket.listenerCount()).toBe(0); + }); + + it('fails the run on an event with a malformed payload', async () => { + const socket = fakeSocket(); + const pending = collect(socketRun(socket, message(), new AbortController().signal, tracking)); + socket.emit({ conversationId: 'c1', type: 'message.delta', payload: { messageId: 'm1' } }); + // Told in words for the user; the field stays on the error for debugging. + await expect(pending).rejects.toThrow('The assistant sent a reply this panel could not read. Try again.'); + await expect(pending).rejects.toMatchObject({ detail: 'Assistant event field "delta" is required' }); + }); + + it('ignores event types the server does not define', () => { + expect( + mapSocketEvent({ + conversationId: 'c', + type: 'proposal.created', + payload: { kind: 'insert', text: 'x' } + } as unknown as AssistantSocketEvent) + ).toBeUndefined(); + expect( + mapSocketEvent({ + conversationId: 'c', + type: 'ui.reveal', + payload: { paragraphId: 'p1' } + } as unknown as AssistantSocketEvent) + ).toBeUndefined(); + }); + + it('drops late responses and aborts runs once a scoped transport is deactivated', async () => { + let resolveList: (conversations: []) => void = () => undefined; + let runSignal: AbortSignal | undefined; + const inner = { + listConversations: vi.fn(() => new Promise<[]>(resolve => (resolveList = resolve))), + createConversation: vi.fn(), + deleteConversation: vi.fn(), + getMessages: vi.fn(), + openRun: async function* (_id: string, _body: unknown, signal: AbortSignal): AsyncIterable { + runSignal = signal; + yield { type: 'run.started' }; + await new Promise(() => undefined); + } + }; + const scoped = scopeTransport(inner); + const list = scoped.transport.listConversations(); + const run = iterate(scoped.transport.openRun('t', { prompt: 'hi' }, new AbortController().signal)); + await run.next(); + expect(inner.listConversations).toHaveBeenCalledWith(); + + scoped.setActive(false); + resolveList([]); + + await expect(list).rejects.toMatchObject({ name: 'AbortError' }); + expect(runSignal?.aborted).toBe(true); + await expect(scoped.transport.getMessages('t')).rejects.toMatchObject({ name: 'AbortError' }); + }); + + it('retires pending creations and active runs when the host session aborts', async () => { + const session = new AbortController(); + let resolveCreate: (value: { id: string }) => void = () => undefined; + let runSignal: AbortSignal | undefined; + const inner = { + listConversations: vi.fn().mockResolvedValue([]), + createConversation: vi.fn( + () => + new Promise<{ id: string }>(resolve => { + resolveCreate = resolve; + }) + ), + deleteConversation: vi.fn(), + getMessages: vi.fn(), + openRun: async function* (_id: string, _body: unknown, signal: AbortSignal): AsyncIterable { + runSignal = signal; + yield { type: 'run.started' }; + await new Promise(resolve => signal.addEventListener('abort', () => resolve(), { once: true })); + yield { type: 'run.completed' }; + } + }; + const scoped = scopeTransport(inner, session.signal); + scoped.setActive(false); + scoped.setActive(true); + const creation = scoped.transport.createConversation({}); + const run = iterate(scoped.transport.openRun('c1', { prompt: 'hi' }, new AbortController().signal)); + await run.next(); + const terminal = run.next(); + session.abort(); + resolveCreate({ id: 'old-account-conversation' }); + await expect(creation).rejects.toMatchObject({ name: 'AbortError' }); + await expect(terminal).rejects.toMatchObject({ name: 'AbortError' }); + expect(runSignal?.aborted).toBe(true); + expect(scoped.active).toBe(false); + scoped.setActive(true); + expect(scoped.active).toBe(false); + await expect(scoped.transport.listConversations()).rejects.toMatchObject({ name: 'AbortError' }); + expect(inner.listConversations).not.toHaveBeenCalled(); + }); + + it('shows one bubble per turn and hides tool and system messages', () => { + expect( + toVisibleMessages([ + { id: 'u1', role: 'user', content: 'improve' }, + { id: 'a1', role: 'assistant', content: '' }, + { id: 't1', role: 'tool', content: '{}' }, + { id: 'a2', role: 'assistant', content: 'Reading.' }, + { id: 't2', role: 'tool', content: '{}' }, + { id: 'a3', role: 'assistant', content: 'Done.' }, + { id: 's1', role: 'system', content: 'hidden' } + ]) + ).toEqual([ + { id: 'u1', role: 'user', content: 'improve' }, + // The merged bubble keeps the id of its last message. + { id: 'a3', role: 'assistant', content: 'Reading.\n\nDone.' } + ]); + }); + + it('keeps the tool calls of a turn on its answer, including across pages', () => { + const visible = toVisibleMessages([ + { id: 'u1', role: 'user', content: 'what is here?' }, + { id: 'a1', role: 'assistant', content: '', toolCalls: [{ id: 'c1', name: 'list_paragraphs' }] }, + { id: 't1', role: 'tool', content: '[]' }, + { id: 'a2', role: 'assistant', content: 'Two paragraphs.' }, + { id: 'u2', role: 'user', content: 'thanks' }, + { id: 'a3', role: 'assistant', content: 'You are welcome.' } + ]); + expect(visible[1]).toEqual({ + id: 'a2', + role: 'assistant', + content: 'Two paragraphs.', + toolCalls: [{ id: 'c1', name: 'list_paragraphs' }] + }); + expect(visible[3]).not.toHaveProperty('toolCalls'); + }); + + it('keeps tool-only messages until an answer arrives, including across raw history pages', () => { + const question = { id: 'u1', role: 'user' as const, content: 'Read this notebook' }; + const call = { + id: 'a1', + role: 'assistant' as const, + content: null, + toolCalls: [{ id: 'c1', name: 'list_paragraphs' }] + }; + const result = { id: 't1', role: 'tool' as const, content: '[]' }; + const answer = { id: 'a2', role: 'assistant' as const, content: 'Two paragraphs.' }; + const earlier = toVisibleMessages([question, call]); + expect(earlier[1]).toEqual({ id: 'a1', role: 'assistant', content: '', toolCalls: call.toolCalls }); + expect(joinPages(earlier, toVisibleMessages([result, answer]))).toEqual( + toVisibleMessages([question, call, result, answer]) + ); + // An intervening page containing only tool results must not lose the pending calls either. + expect(joinPages(joinPages(earlier, toVisibleMessages([result])), toVisibleMessages([answer]))).toEqual([ + question, + { ...answer, toolCalls: call.toolCalls } + ]); + const nextTurn = [ + { id: 'u2', role: 'user' as const, content: 'Another question' }, + { id: 'a3', role: 'assistant' as const, content: 'Another answer' } + ]; + expect(joinPages(earlier, toVisibleMessages(nextTurn))).toEqual([...earlier, ...nextTurn]); + expect(toVisibleMessages([question, call, ...nextTurn])).toEqual([...earlier, ...nextTurn]); + }); + + it('keeps a trailing tool round after answer text and deduplicates call ids at page boundaries', () => { + const call = { id: 'c1', name: 'list_paragraphs' }; + const earlier = toVisibleMessages([ + { id: 'a1', role: 'assistant', content: 'Looking.' }, + { id: 'a2', role: 'assistant', content: '', toolCalls: [call] } + ]); + expect(earlier).toEqual([{ id: 'a2', role: 'assistant', content: 'Looking.', toolCalls: [call] }]); + const later = toVisibleMessages([ + { id: 'a2', role: 'assistant', content: '', toolCalls: [call] }, + { id: 'a3', role: 'assistant', content: 'Done.' } + ]); + expect(joinPages(earlier, later)).toEqual([ + { id: 'a3', role: 'assistant', content: 'Looking.\n\nDone.', toolCalls: [call] } + ]); + }); +}); + +describe('retired run subscriptions', () => { + it('does not subscribe or throw for an inactive or aborted session', () => { + const inner = createAssistantTransport('/api', 'note', fakeSocket()); + const subscribe = vi.spyOn(inner, 'subscribeRunState'); + const signal = new AbortController(); + const scoped = scopeTransport(inner, signal.signal); + signal.abort(); + expect(() => scoped.transport.subscribeRunState!('conversation', vi.fn())()).not.toThrow(); + expect(subscribe).not.toHaveBeenCalled(); + const retired = scopeTransport(inner); + retired.setActive(false); + expect(() => retired.transport.subscribeRunState!('conversation', vi.fn())()).not.toThrow(); + expect(subscribe).not.toHaveBeenCalled(); + }); +}); + +describe('retained run state', () => { + it.each([true, false])( + 'keeps old completion separate from a new rejection (completion first: %s)', + async completionFirst => { + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const controller = new AbortController(); + const first = iterate(transport.openRun('conversation', { prompt: 'old' }, controller.signal)); + const started = first.next(); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'old' } }); + await started; + controller.abort(); + await expect(first.next()).rejects.toThrow(); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', running => states.push(running)); + const next = collect(transport.openRun('conversation', { prompt: 'new' }, new AbortController().signal)); + const completion = () => + socket.emit({ conversationId: 'conversation', type: 'run.completed', payload: { runId: 'old' } }); + const rejection = () => + socket.emit({ + conversationId: 'conversation', + type: 'run.failed', + payload: { runId: 'new', error: { status: 429 } } + }); + if (completionFirst) { + completion(); + rejection(); + } else { + rejection(); + completion(); + } + expect(await next).toEqual([ + { type: 'run.failed', runId: 'new', message: 'The assistant is busy. Try again shortly.' } + ]); + expect(states[0]).toBe('running'); + expect(states[states.length - 1]).toBe('idle'); + expect(socket.listenerCount()).toBe(0); + stop(); + } + ); + + it('uses the server conversation response without a running field', async () => { + const fetch = vi + .fn() + .mockResolvedValue(new Response(JSON.stringify({ body: [{ id: 'conversation', title: 'Question' }] }))); + vi.stubGlobal('fetch', fetch); + const transport = createAssistantTransport('/api', 'note', fakeSocket()); + expect(await transport.listConversations()).toEqual([{ id: 'conversation', title: 'Question' }]); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', state => states.push(state)); + expect(states).toEqual(['idle']); + expect(fetch).toHaveBeenCalledTimes(1); + stop(); + }); + + it('keeps a disconnected run visible until history reload succeeds', async () => { + const fetch = vi + .fn() + .mockRejectedValueOnce(new TypeError('Offline')) + .mockResolvedValueOnce(new Response(JSON.stringify({ body: { messages: [], cursor: null } }))) + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + body: { messages: [{ id: 'answer', role: 'assistant', content: 'Recovered answer' }], cursor: null } + }) + ) + ); + vi.stubGlobal('fetch', fetch); + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', state => states.push(state)); + const iterator = iterate(transport.openRun('conversation', { prompt: 'question' }, new AbortController().signal)); + const first = iterator.next(); + socket.close(); + await expect(first).rejects.toBeInstanceOf(AssistantConnectionError); + expect(states[states.length - 1]).toBe('disconnected'); + await expect(transport.getMessages('conversation')).rejects.toThrow('Offline'); + expect(states[states.length - 1]).toBe('disconnected'); + await transport.getMessages('conversation', 'older-page'); + expect(states[states.length - 1]).toBe('disconnected'); + await expect(transport.getMessages('conversation')).resolves.toEqual({ + messages: [{ id: 'answer', role: 'assistant', content: 'Recovered answer' }], + earlierCursor: null + }); + expect(states[states.length - 1]).toBe('idle'); + expect(fetch).toHaveBeenCalledTimes(3); + stop(); + }); + + it('keeps a disconnected run visible when a retry is rejected', async () => { + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const first = collect(transport.openRun('conversation', { prompt: 'first' }, new AbortController().signal)); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'first' } }); + socket.close(); + await expect(first).rejects.toBeInstanceOf(AssistantConnectionError); + + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', state => states.push(state)); + const retry = collect(transport.openRun('conversation', { prompt: 'retry' }, new AbortController().signal)); + socket.emit({ + conversationId: 'conversation', + type: 'run.failed', + payload: { runId: 'retry', error: { status: 409 } } + }); + await expect(retry).resolves.toEqual([ + { + type: 'run.failed', + runId: 'retry', + message: 'This conversation is still answering. Try again when it finishes.' + } + ]); + expect(states).toContain('disconnected'); + expect(states).not.toContain('idle'); + + const accepted = collect(transport.openRun('conversation', { prompt: 'later' }, new AbortController().signal)); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'later' } }); + socket.emit({ conversationId: 'conversation', type: 'run.completed', payload: { runId: 'later' } }); + await expect(accepted).resolves.toEqual([ + { type: 'run.started', runId: 'later' }, + { type: 'run.completed', runId: 'later' } + ]); + expect(states).toContain('running'); + expect(states.at(-1)).toBe('idle'); + stop(); + }); + + it('does not recover a disconnect with history requested before it', async () => { + let resolveHistory: (response: Response) => void = () => undefined; + const delayedHistory = new Promise(resolve => { + resolveHistory = resolve; + }); + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockReturnValueOnce(delayedHistory) + .mockResolvedValueOnce(new Response(JSON.stringify({ body: { messages: [], cursor: null } }))) + ); + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const run = collect(transport.openRun('conversation', { prompt: 'question' }, new AbortController().signal)); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'run' } }); + const staleHistory = transport.getMessages('conversation'); + socket.close(); + await expect(run).rejects.toBeInstanceOf(AssistantConnectionError); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', state => states.push(state)); + + resolveHistory(new Response(JSON.stringify({ body: { messages: [], cursor: null } }))); + await staleHistory; + expect(states).toEqual(['disconnected']); + + await transport.getMessages('conversation'); + expect(states).toEqual(['disconnected', 'idle']); + stop(); + }); + + it('clears a retained run on disconnect and releases listeners', async () => { + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const controller = new AbortController(); + const iterator = iterate(transport.openRun('conversation', { prompt: 'question' }, controller.signal)); + const first = iterator.next(); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'run' } }); + await first; + controller.abort(); + await expect(iterator.next()).rejects.toThrow(); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', running => states.push(running)); + socket.close(); + expect(states).toEqual(['running', 'disconnected']); + expect(socket.listenerCount()).toBe(0); + expect(socket.closeListenerCount()).toBe(0); + stop(); + }); + + it('follows an abandoned accepted run until completion and releases socket listeners', async () => { + const socket = fakeSocket(); + const transport = createAssistantTransport('/api', 'note', socket); + const states: AssistantRunState[] = []; + const stop = transport.subscribeRunState!('conversation', running => states.push(running)); + const controller = new AbortController(); + const iterator = iterate(transport.openRun('conversation', { prompt: 'question' }, controller.signal)); + const first = iterator.next(); + socket.emit({ conversationId: 'conversation', type: 'run.started', payload: { runId: 'run' } }); + expect((await first).value).toEqual({ type: 'run.started', runId: 'run' }); + controller.abort(); + await expect(iterator.next()).rejects.toThrow(); + const reopened: AssistantRunState[] = []; + const stopReopened = transport.subscribeRunState!('conversation', running => reopened.push(running)); + expect(reopened).toEqual(['running']); + socket.emit({ conversationId: 'conversation', type: 'run.completed', payload: { runId: 'run' } }); + expect(reopened).toEqual(['running', 'idle']); + expect(states).toContain('running'); + expect(socket.listenerCount()).toBe(0); + expect(socket.closeListenerCount()).toBe(0); + stop(); + stopReopened(); + }); +}); diff --git a/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.ts b/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.ts new file mode 100644 index 00000000000..65211ac93cb --- /dev/null +++ b/zeppelin-web-angular/projects/zeppelin-react/src/entities/assistant/model/assistantTransport.ts @@ -0,0 +1,599 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { AssistantSendMessage, AssistantSocket, AssistantSocketEvent } from '@zeppelin/sdk'; +import { + AssistantConnectionError, + AssistantStreamError, + type AssistantConversation, + type AssistantMessage, + type AssistantMessagePage, + type AssistantRunEvent, + type AssistantRunState, + type AssistantToolCallRef, + type AssistantTransport +} from './assistantContract'; +import { mergeAnswers } from './messageHistory'; + +const HTTP_MESSAGES: Record = { + 403: 'You do not have permission for this conversation.', + 404: 'This conversation or note no longer exists.', + 409: 'This conversation is still answering. Try again when it finishes.', + 503: 'The assistant is not configured on this server. Ask an administrator to enable it.' +}; + +// run.failed carries only the HTTP status of the rejected or failed run, so the UI words it. +const RUN_ERROR_MESSAGES: Record = { + ...HTTP_MESSAGES, + 400: 'The assistant could not accept this message. Edit it and try again.', + 429: 'The assistant is busy. Try again shortly.', + // Also returned when the user lost read access to the note, so it does not name the owner. + 403: 'You do not have permission to send messages to this conversation.' +}; +const NOT_AVAILABLE_MESSAGE = 'The assistant is not available on this server, or this note no longer exists.'; +const RUN_ERROR_FALLBACK = 'The assistant hit an error while answering. Try again.'; + +export class AssistantHttpError extends Error { + constructor( + public readonly status: number, + public readonly location: string | null, + message = HTTP_MESSAGES[status] ?? `Assistant request failed with HTTP ${status}` + ) { + super(message); + this.name = 'AssistantHttpError'; + } +} + +/** An event this panel cannot read; the field stays on the error for debugging, the message is for the user. */ +export class AssistantProtocolError extends Error { + constructor(readonly detail: string) { + super('The assistant sent a reply this panel could not read. Try again.'); + this.name = 'AssistantProtocolError'; + } +} + +// The server sends no heartbeat, so a run that goes quiet this long is treated as lost. +const RUN_IDLE_TIMEOUT_MS = 180_000; + +const isRecord = (value: unknown): value is Record => typeof value === 'object' && value !== null; + +const optionalString = (data: Record, field: string): string | undefined => { + const value = data[field]; + if (value !== undefined && typeof value !== 'string') { + throw new AssistantProtocolError(`Assistant event field "${field}" must be a string`); + } + return value; +}; + +const requiredString = (data: Record, field: string): string => { + const value = optionalString(data, field); + if (value === undefined) { + throw new AssistantProtocolError(`Assistant event field "${field}" is required`); + } + return value; +}; + +const errorMessage = (data: Record): string => { + const error = data.error; + const status = isRecord(error) ? error.status : undefined; + return (typeof status === 'number' && RUN_ERROR_MESSAGES[status]) || RUN_ERROR_FALLBACK; +}; + +/** Maps an ASSISTANT_EVENT to the UI contract; unknown types map to undefined, so new server events are ignored. */ +export const mapSocketEvent = ({ type, payload }: AssistantSocketEvent): AssistantRunEvent | undefined => { + const data = isRecord(payload) ? payload : {}; + switch (type) { + case 'run.started': + return { type, runId: optionalString(data, 'runId') }; + case 'run.completed': + return { type, runId: optionalString(data, 'runId') }; + case 'run.failed': + return { type, runId: optionalString(data, 'runId'), message: errorMessage(data) }; + case 'message.delta': + return { type, messageId: requiredString(data, 'messageId'), delta: requiredString(data, 'delta') }; + case 'message.done': + return { type, messageId: requiredString(data, 'messageId'), content: requiredString(data, 'content') }; + case 'tool_call.started': + return { type, toolCallId: requiredString(data, 'toolCallId'), name: requiredString(data, 'name') }; + case 'tool_call.done': + return { type, toolCallId: requiredString(data, 'toolCallId'), name: optionalString(data, 'name') }; + default: + return undefined; + } +}; + +const isTerminal = (event: AssistantRunEvent): boolean => event.type === 'run.completed' || event.type === 'run.failed'; + +interface ConversationMessage { + id: string; + role: 'user' | 'assistant' | 'tool' | 'system'; + content?: string | null; + toolCalls?: Array<{ id?: unknown; name?: unknown }> | null; +} + +const toolCallRefs = (message: ConversationMessage): AssistantToolCallRef[] => + (message.toolCalls ?? []).flatMap(call => + typeof call?.id === 'string' && typeof call.name === 'string' ? [{ id: call.id, name: call.name }] : [] + ); + +/** + * One bubble per turn: tool and system messages are hidden, tool rounds merge under the last message's id, and + * the tools called on the way are kept on the answer so a reopened conversation still shows its action log. + */ +export const toVisibleMessages = (messages: ConversationMessage[]): AssistantMessage[] => { + const visible: AssistantMessage[] = []; + for (const message of messages) { + if (message.role === 'user') { + visible.push({ id: message.id, role: 'user', content: message.content ?? '' }); + } else if (message.role === 'assistant') { + const toolCalls = toolCallRefs(message); + // Keep a tool-only turn even when its answer is on a later page or has not arrived yet. + if (!message.content && !toolCalls.length) continue; + const last = visible[visible.length - 1]; + const answer: AssistantMessage = { + id: message.id, + role: 'assistant', + content: message.content ?? '', + ...(toolCalls.length ? { toolCalls } : {}) + }; + if (last?.role === 'assistant') visible[visible.length - 1] = mergeAnswers(last, answer); + else visible.push(answer); + } + } + return visible; +}; + +const requestHeaders = { + // As the shell does in production, so the server answers 401/405 instead of a login page. + 'X-Requested-With': 'XMLHttpRequest' +}; + +// Zeppelin JSON endpoints wrap the payload as { status, message, body }. +const unwrapBody = (value: unknown): unknown => + typeof value === 'object' && value !== null && 'body' in value ? value.body : value; + +/** Hands 401/405 to the host, which owns login redirects and logout. */ +export type AssistantAuthErrorHandler = (status: number, location: string | null) => void; + +const httpError = (response: Response, onAuthError?: AssistantAuthErrorHandler): AssistantHttpError => { + const error = new AssistantHttpError(response.status, response.headers.get('Location')); + if (error.status === 401 || error.status === 405) { + onAuthError?.(error.status, error.location); + } + return error; +}; + +const requestJson = async ( + url: string, + init: { method?: string; body?: unknown; signal?: AbortSignal } = {}, + onAuthError?: AssistantAuthErrorHandler +): Promise => { + const response = await fetch(url, { + method: init.method ?? 'GET', + ...(init.signal ? { signal: init.signal } : {}), + credentials: 'include', + headers: + init.body === undefined + ? { ...requestHeaders, Accept: 'application/json' } + : { ...requestHeaders, Accept: 'application/json', 'Content-Type': 'application/json' }, + body: init.body === undefined ? undefined : JSON.stringify(init.body) + }); + if (!response.ok) { + throw httpError(response, onAuthError); + } + const text = await response.text(); + return (text ? unwrapBody(JSON.parse(text)) : undefined) as T; +}; + +// A server without the assistant API answers the note's conversation collection with 404. +const collectionRequest = async (...args: Parameters): Promise => { + try { + return await requestJson(...args); + } catch (error) { + if (error instanceof AssistantHttpError && error.status === 404) { + throw new AssistantHttpError(404, error.location, NOT_AVAILABLE_MESSAGE); + } + throw error; + } +}; + +const MAX_ABANDONED_RUNS = 200; +/** + * Runs on one host connection. Pending sends survive note switches until acknowledged; abandoned run ids keep + * later sends from following their events. The server continues runs after the UI stops listening. + */ +interface PendingRun { + abandoned: boolean; + runId?: string; +} +export interface RunTracking { + abandoned: Set; + pending: Map; + active: Map; + disconnected: Map; + listeners: Map void>>; + stopTracking?: () => void; +} +export const createRunTracking = (): RunTracking => ({ + abandoned: new Set(), + pending: new Map(), + active: new Map(), + disconnected: new Map(), + listeners: new Map() +}); +const runState = (tracking: RunTracking, conversationId: string): AssistantRunState => + tracking.disconnected.has(conversationId) + ? 'disconnected' + : tracking.pending.has(conversationId) || tracking.active.has(conversationId) + ? 'running' + : 'idle'; +const notifyRunState = (tracking: RunTracking, conversationId: string) => { + tracking.listeners.get(conversationId)?.forEach(listener => listener(runState(tracking, conversationId))); +}; +const releaseTracking = (tracking: RunTracking) => { + if (!tracking.pending.size && !tracking.active.size) tracking.stopTracking?.(); +}; +const connectionRuns = new WeakMap(); +const runsFor = (socket: AssistantSocket): RunTracking => { + const key = socket.connectionKey ?? socket; + let runs = connectionRuns.get(key); + if (!runs) connectionRuns.set(key, (runs = createRunTracking())); + return runs; +}; +const abandon = (tracking: RunTracking, runId: string) => { + tracking.abandoned.add(runId); + if (tracking.abandoned.size > MAX_ABANDONED_RUNS) { + const oldest = tracking.abandoned.values().next().value; + if (oldest !== undefined) tracking.abandoned.delete(oldest); + } +}; + +// ACKs belong to the connection, even if the note or conversation that sent the request has unmounted. +const trackPendingRun = (socket: AssistantSocket, conversationId: string, tracking: RunTracking): PendingRun => { + const pending: PendingRun = { abandoned: false }; + tracking.pending.set(conversationId, pending); + notifyRunState(tracking, conversationId); + if (!tracking.stopTracking) { + const unsubscribe = socket.subscribe(event => { + const request = tracking.pending.get(event.conversationId); + const payload = isRecord(event.payload) ? event.payload : {}; + const runId = typeof payload.runId === 'string' ? payload.runId : undefined; + if (event.type === 'run.started' && runId && request) { + request.runId = runId; + tracking.active.set(event.conversationId, runId); + tracking.disconnected.delete(event.conversationId); + if (request.abandoned) abandon(tracking, runId); + tracking.pending.delete(event.conversationId); + } else if (event.type === 'run.failed' || event.type === 'run.completed') { + const activeRun = tracking.active.get(event.conversationId); + if (activeRun && (!runId || runId === activeRun)) { + tracking.active.delete(event.conversationId); + } else if (request && !(runId && tracking.abandoned.has(runId))) { + tracking.pending.delete(event.conversationId); + } else return; + } else return; + notifyRunState(tracking, event.conversationId); + releaseTracking(tracking); + }); + const unsubscribeClose = socket.subscribeClose(() => { + const conversations = new Set([...Array.from(tracking.pending.keys()), ...Array.from(tracking.active.keys())]); + tracking.pending.clear(); + tracking.active.clear(); + tracking.abandoned.clear(); + conversations.forEach(id => tracking.disconnected.set(id, {})); + conversations.forEach(id => notifyRunState(tracking, id)); + tracking.stopTracking?.(); + }); + tracking.stopTracking = () => { + unsubscribe(); + unsubscribeClose(); + tracking.stopTracking = undefined; + }; + } + return pending; +}; + +/** + * Sends one message over the host's notebook WebSocket and yields that conversation's events until a terminal one. + * Until this message's run.started arrives, only a rejection (run.failed) belongs to it; events of an abandoned run + * on the same conversation are skipped. + */ +export async function* socketRun( + socket: AssistantSocket, + message: AssistantSendMessage, + signal: AbortSignal, + tracking: RunTracking, + idleTimeoutMs = RUN_IDLE_TIMEOUT_MS +): AsyncIterable { + signal.throwIfAborted(); + const { abandoned } = tracking; + if (tracking.pending.has(message.conversationId)) { + yield { type: 'run.failed', message: 'The previous request is still being confirmed. Try again shortly.' }; + return; + } + const pending = trackPendingRun(socket, message.conversationId, tracking); + let sent = false; + const queue: Array = []; + let wake: (() => void) | undefined; + const push = (item: AssistantRunEvent | Error) => { + queue.push(item); + wake?.(); + }; + const unsubscribe = socket.subscribe(event => { + if (event.conversationId !== message.conversationId) return; + try { + const mapped = mapSocketEvent(event); + if (mapped) push(mapped); + } catch (error) { + push(error instanceof Error ? error : new Error(String(error))); + } + }); + const unsubscribeClose = socket.subscribeClose(() => push(new AssistantConnectionError())); + const onAbort = () => { + pending.abandoned = true; + if (pending.runId) abandon(tracking, pending.runId); + wake?.(); + }; + signal.addEventListener('abort', onAbort, { once: true }); + let started = false; + let ownRunId: string | undefined; + let ended = false; + try { + socket.send(message); + sent = true; + while (true) { + signal.throwIfAborted(); + if (queue.length === 0) { + let timer: ReturnType | undefined; + await new Promise(resolve => { + wake = resolve; + timer = setTimeout(() => push(new AssistantStreamError()), idleTimeoutMs); + }); + clearTimeout(timer); + wake = undefined; + } + signal.throwIfAborted(); + const item = queue.shift(); + if (item === undefined) continue; + if (item instanceof Error) throw item; + const runId = 'runId' in item ? item.runId : undefined; + if (runId && runId !== ownRunId && abandoned.has(runId)) { + if (isTerminal(item)) abandoned.delete(runId); + continue; + } + // Once running, only its own terminal event ends it; deltas have no run id, but runs are one at a time. + if (started && isTerminal(item) && runId && ownRunId && runId !== ownRunId) continue; + if (!started) { + if (item.type === 'run.started') { + started = true; + ownRunId = item.runId; + } else if (item.type !== 'run.failed') continue; + } + if (isTerminal(item)) ended = true; + yield item; + if (ended) return; + } + } finally { + if (!ended && sent) { + pending.abandoned = true; + const runId = pending.runId ?? ownRunId; + if (runId) abandon(tracking, runId); + } + if (!sent || ended) { + if (tracking.pending.get(message.conversationId) === pending) tracking.pending.delete(message.conversationId); + notifyRunState(tracking, message.conversationId); + releaseTracking(tracking); + } + signal.removeEventListener('abort', onAbort); + unsubscribe(); + unsubscribeClose(); + } +} + +type ConversationSummary = { + id: string; + title?: string; + ownerId?: string; + canSendMessage?: boolean; + createdAt?: string; + updatedAt?: string; +}; +type MessagePage = { messages: ConversationMessage[]; cursor?: string | null }; + +// One page of raw server messages; the panel loads earlier pages as the user scrolls up. +const HISTORY_PAGE_SIZE = 50; + +/** + * Transport for the per-note conversation API. REST for conversations and history (`apiBase` is the host's + * REST base, e.g. `.../api`); sending goes over the host's notebook WebSocket. + */ +export const createAssistantTransport = ( + apiBase: string, + noteId: string, + socket: AssistantSocket, + onAuthError?: AssistantAuthErrorHandler +): AssistantTransport => { + const base = `${apiBase.replace(/\/$/, '')}/notes/${encodeURIComponent(noteId)}/conversations`; + const runs = runsFor(socket); + const toConversation = (conversation: ConversationSummary): AssistantConversation => { + return { + id: conversation.id, + title: conversation.title, + ...(conversation.ownerId !== undefined ? { ownerId: conversation.ownerId } : {}), + ...(typeof conversation.canSendMessage === 'boolean' ? { canSendMessage: conversation.canSendMessage } : {}), + ...((conversation.updatedAt ?? conversation.createdAt) + ? { updatedAt: conversation.updatedAt ?? conversation.createdAt } + : {}) + }; + }; + // The server names an untitled conversation by its creation time. + const titleFor = (title?: string) => { + const trimmed = title?.trim().replace(/\s+/g, ' '); + return trimmed ? { title: trimmed.length > 80 ? `${trimmed.slice(0, 79)}…` : trimmed } : {}; + }; + return { + subscribeRunState: (conversationId, listener) => { + let listeners = runs.listeners.get(conversationId); + if (!listeners) runs.listeners.set(conversationId, (listeners = new Set())); + let stopped = false; + const observe = (state: AssistantRunState) => { + listener(state); + }; + listeners.add(observe); + observe(runState(runs, conversationId)); + const stop = () => { + if (stopped) return; + stopped = true; + listeners.delete(observe); + if (!listeners.size) runs.listeners.delete(conversationId); + socket.signal?.removeEventListener('abort', stop); + }; + socket.signal?.addEventListener('abort', stop, { once: true }); + if (socket.signal?.aborted) stop(); + return stop; + }, + // Latest activity first, so the panel opens on it; undated ones stay last in server order (stable sort). + listConversations: async () => + (await collectionRequest(base, {}, onAuthError)) + .map(conversation => ({ + conversation, + time: Date.parse(conversation.updatedAt ?? conversation.createdAt ?? '') || 0 + })) + .sort((a, b) => b.time - a.time) + .map(({ conversation }) => toConversation(conversation)), + createConversation: async ({ title } = {}) => + toConversation( + await collectionRequest(base, { method: 'POST', body: titleFor(title) }, onAuthError) + ), + deleteConversation: async conversationId => { + await requestJson(`${base}/${encodeURIComponent(conversationId)}`, { method: 'DELETE' }, onAuthError); + }, + getMessages: async (conversationId, before): Promise => { + const cursor = before ? `&cursor=${encodeURIComponent(before)}` : ''; + const disconnectedAtRequest = runs.disconnected.get(conversationId); + const page = await requestJson( + `${base}/${encodeURIComponent(conversationId)}/messages?limit=${HISTORY_PAGE_SIZE}${cursor}`, + {}, + onAuthError + ); + if (!before && disconnectedAtRequest && runs.disconnected.get(conversationId) === disconnectedAtRequest) { + runs.disconnected.delete(conversationId); + notifyRunState(runs, conversationId); + } + // The server pages from the latest message backwards; the panel renders oldest to newest. + return { messages: toVisibleMessages([...page.messages].reverse()), earlierCursor: page.cursor ?? null }; + }, + openRun: (conversationId, body, signal) => + socketRun( + socket, + { + noteId, + conversationId, + content: body.prompt + }, + signal, + runs + ) + }; +}; + +/** + * Guards a per-note transport. Deactivation aborts open runs and rejects calls while inactive. + * The transport can be reactivated unless its session signal has aborted. + */ +export const scopeTransport = (inner: AssistantTransport, sessionSignal?: AbortSignal) => { + let active = true; + const runs = new Set(); + const subscriptions = new Set<() => void>(); + const assertActive = () => { + if (!active || sessionSignal?.aborted) { + throw new DOMException('Notebook changed', 'AbortError'); + } + }; + const guarded = async (request: () => Promise): Promise => { + assertActive(); + try { + const result = await request(); + assertActive(); + return result; + } catch (error) { + assertActive(); + throw error; + } + }; + const transport: AssistantTransport = { + listConversations: () => guarded(() => inner.listConversations()), + createConversation: body => guarded(() => inner.createConversation(body)), + deleteConversation: id => guarded(() => inner.deleteConversation(id)), + getMessages: (id, before) => guarded(() => inner.getMessages(id, before)), + ...(inner.subscribeRunState + ? { + subscribeRunState: (id: string, listener: (state: AssistantRunState) => void) => { + if (!active || sessionSignal?.aborted) return () => undefined; + const unsubscribe = inner.subscribeRunState!(id, state => { + if (active && !sessionSignal?.aborted) listener(state); + }); + const stop = () => { + if (!subscriptions.delete(stop)) return; + unsubscribe(); + sessionSignal?.removeEventListener('abort', stop); + }; + subscriptions.add(stop); + sessionSignal?.addEventListener('abort', stop, { once: true }); + return stop; + } + } + : {}), + async *openRun(id, body, signal) { + assertActive(); + const controller = new AbortController(); + const abort = () => controller.abort(); + signal.addEventListener('abort', abort, { once: true }); + sessionSignal?.addEventListener('abort', abort, { once: true }); + if (signal.aborted) { + abort(); + } + runs.add(controller); + try { + for await (const event of inner.openRun(id, body, controller.signal)) { + assertActive(); + if (controller.signal.aborted) { + throw new DOMException('Run stopped', 'AbortError'); + } + yield event; + } + } catch (error) { + assertActive(); + throw error; + } finally { + controller.abort(); + signal.removeEventListener('abort', abort); + sessionSignal?.removeEventListener('abort', abort); + runs.delete(controller); + } + } + }; + return { + get active() { + return active && !sessionSignal?.aborted; + }, + transport, + setActive(next: boolean) { + active = next; + if (!next) { + runs.forEach(controller => controller.abort()); + runs.clear(); + subscriptions.forEach(stop => stop()); + } + } + }; +}; diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/assistant.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/assistant.ts new file mode 100644 index 00000000000..13324721ca1 --- /dev/null +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/assistant.ts @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { AssistantParagraphRef, AssistantRevealResult, AssistantSlot } from './assistant-ui'; + +import type { + AssistantEvent as AssistantSocketEvent, + AssistantSendMessage +} from './interfaces/message-assistant.interface'; + +/** ASSISTANT_EVENT as the server sends it; the remote maps it to its run events. */ +export type { AssistantSocketEvent }; + +/** The host's notebook WebSocket, narrowed to the assistant ops. */ +export interface AssistantSocket { + /** Stable host connection identity across notebook adapters. */ + readonly connectionKey?: object; + /** Aborted when the host retires this notebook/account adapter, without closing the shared connection. */ + readonly signal?: AbortSignal; + send(message: AssistantSendMessage): void; + subscribe(listener: (event: AssistantSocketEvent) => void): () => void; + /** The connection closed. The server sends a run's events only to the connection that started it. */ + subscribeClose(listener: () => void): () => void; +} + +/** Props the Angular notebook passes to the `./AssistantWorkspace` remote. A type alias, so it fits `ReactProps`. */ +export type AssistantHostProps = { + noteId: string; + /** The host's REST base (`.../api`); the remote calls the conversation REST API itself. */ + apiBase: string; + /** The host's notebook WebSocket, used to send messages and receive run events. */ + socket: AssistantSocket; + draftOwner?: string; + slots: AssistantSlot[]; + onAuthError?: (status: number, location: string | null) => void; + onError?: (error: unknown) => void; + /** Keeps the host sidebar to one open view at a time. */ + onPanelVisibilityChange?: (visible: boolean) => void; + subscribePanelClose?: (listener: () => void) => () => void; + /** Host-side scroll and highlight for a paragraph an answer links to; notebook DOM stays with Angular. */ + revealParagraph?: (paragraphId: string) => Promise; + /** The notebook's paragraphs in order, as data so labelling them never calls into the host. */ + paragraphs?: AssistantParagraphRef[]; + /** The notebook sidebar's width, which the panel shares so switching views keeps it. */ + panelWidth?: number; + /** A resize of the panel, at the end of a drag or per arrow key. */ + onPanelWidthChange?: (width: number) => void; +}; diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/public-api.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/public-api.ts index dcf432e8446..2095d38fe9c 100644 --- a/zeppelin-web-angular/projects/zeppelin-sdk/src/public-api.ts +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/public-api.ts @@ -11,8 +11,9 @@ */ export * from './angular-object'; +export * from './assistant'; export * from './assistant-ui'; export * from './interfaces/public-api'; -export * from './message'; +export * from './message'; export * from './paragraph-state';