GPU_GUARD_MONOREPO/packages/backend/test/chat_orchestrator.test.ts

jest.mock('uuid', () => ({
  v4: jest.fn(() => 'session-id'),
}));

import { NetaClawSessionService } from '../src/modules/netaclaw/gateway/session.js';
import { NetaClawChatOrchestratorService } from '../src/modules/netaclaw/service/chat_orchestrator.js';

describe('NetaClawChatOrchestratorService', () => {
  it('anchors an assistant placeholder before running the agent and updates it with the final content', async () => {
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([{ id: 101, role: 'user', content: 'hello' }]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'hello' })
        .mockResolvedValueOnce({
          id: 202,
          sessionId: 'session-1',
          role: 'assistant',
          content: '',
          metadata: { inFlight: true },
        }),
      updateMessage: jest.fn().mockResolvedValue(undefined),
    } as unknown as NetaClawSessionService;

    const runner = jest.fn(async params => {
      expect(params.assistantMessageId).toBe(202);
      expect(params.userMessageId).toBe(101);
      expect(params.history).toEqual([]);
      expect(sessionService.saveMessage).toHaveBeenCalledTimes(2);
      expect(sessionService.updateMessage).not.toHaveBeenCalled();

      return {
        finalContent: 'final answer',
        usage: { inputTokens: 11, outputTokens: 7 },
        toolCallCount: 0,
      };
    });

    const service = new NetaClawChatOrchestratorService();
    service.sessionService = sessionService;

    const result = await service.executeChat({
      sessionId: 'session-1',
      requestedAgentName: 'default',
      message: 'hello',
      agentId: 7,
      userId: 'user-1',
      runner,
    });

    expect(sessionService.getOrCreateSession).toHaveBeenCalledWith('session-1', 'default', 7, 'user-1');
    expect(sessionService.saveMessage).toHaveBeenNthCalledWith(1, 'session-1', {
      role: 'user',
      content: 'hello',
    });
    expect(sessionService.saveMessage).toHaveBeenNthCalledWith(2, 'session-1', {
      role: 'assistant',
      content: '',
      metadata: { inFlight: true },
    });
    expect(runner).toHaveBeenCalledWith(
      expect.objectContaining({
        sessionId: 'session-1',
        assistantMessageId: 202,
        userMessageId: 101,
        userMessage: 'hello',
      })
    );
    expect(sessionService.updateMessage).toHaveBeenCalledWith(202, {
      content: 'final answer',
      metadata: { inFlight: false },
    });
    expect(result).toEqual({
      sessionId: 'session-1',
      finalContent: 'final answer',
      content: 'final answer',
      usage: { inputTokens: 11, outputTokens: 7 },
      toolCallCount: 0,
      assistantMessageId: 202,
      subagentBatches: [],
    });
  });

  it('excludes the just-saved user message by exact id even when later history entries exist', async () => {
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([
        { id: 11, role: 'user', content: 'older user' },
        { id: 101, role: 'user', content: 'current user' },
        { id: 150, role: 'assistant', content: 'later assistant' },
      ]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'current user' })
        .mockResolvedValueOnce({
          id: 202,
          sessionId: 'session-1',
          role: 'assistant',
          content: '',
          metadata: { inFlight: true },
        }),
      updateMessage: jest.fn().mockResolvedValue(undefined),
    } as unknown as NetaClawSessionService;

    const runner = jest.fn().mockResolvedValue({
      finalContent: 'final answer',
      toolCallCount: 0,
    });

    const service = new NetaClawChatOrchestratorService();
    service.sessionService = sessionService;

    await service.executeChat({
      sessionId: 'session-1',
      requestedAgentName: 'default',
      message: 'current user',
      runner,
    });

    expect(runner).toHaveBeenCalledWith(
      expect.objectContaining({
        history: [
          { role: 'user', content: 'older user', toolCalls: undefined, toolCallId: undefined },
          { role: 'assistant', content: 'later assistant', toolCalls: undefined, toolCallId: undefined },
        ],
      })
    );
  });

  it('passes the saved user message id into history preparation callbacks', async () => {
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([
        { id: 7, role: 'assistant', content: 'older answer' },
        { id: 101, role: 'user', content: 'current user' },
      ]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'current user' })
        .mockResolvedValueOnce({ id: 202, sessionId: 'session-1', role: 'assistant', content: '' }),
      updateMessage: jest.fn().mockResolvedValue(undefined),
    } as unknown as NetaClawSessionService;

    const prepareHistory = jest.fn(async ({ userMessageId, history }) => {
      expect(userMessageId).toBe(101);
      return history;
    });
    const runner = jest.fn().mockResolvedValue({
      finalContent: 'final answer',
      toolCallCount: 0,
    });

    const service = new NetaClawChatOrchestratorService();
    service.sessionService = sessionService;

    await service.executeChat({
      sessionId: 'session-1',
      requestedAgentName: 'default',
      message: 'current user',
      prepareHistory,
      runner,
    });

    expect(prepareHistory).toHaveBeenCalledWith({
      sessionId: 'session-1',
      userMessageId: 101,
      history: [{ role: 'assistant', content: 'older answer', toolCalls: undefined, toolCallId: undefined }],
    });
  });

  it('collects subagent batches for the assistant placeholder and persists them on the same assistant message', async () => {
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([{ id: 101, role: 'user', content: 'delegate this' }]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'delegate this' })
        .mockResolvedValueOnce({
          id: 202,
          sessionId: 'session-1',
          role: 'assistant',
          content: '',
          metadata: { inFlight: true },
        }),
      updateMessage: jest.fn().mockResolvedValue(undefined),
    } as unknown as NetaClawSessionService;

    const service = new NetaClawChatOrchestratorService();
    service.sessionService = sessionService;

    const result = await service.executeChat({
      sessionId: 'session-1',
      requestedAgentName: 'default',
      message: 'delegate this',
      runner: async params => {
        expect(params.assistantMessageId).toBe(202);
        expect(params.subagents).toEqual(
          expect.objectContaining({
            sessionId: 'session-1',
            parentMessageId: 202,
            batches: [],
          })
        );

        params.subagents.onBatchStart({
          batchId: 'batch-1',
          sessionId: 'session-1',
          parentMessageId: 202,
          mode: 'parallel',
          tasks: [
            { goal: 'task-1', agentName: 'worker-a' },
            { goal: 'task-2', agentName: 'worker-b' },
          ],
        });
        params.subagents.onUpdate({
          batchId: 'batch-1',
          sessionId: 'session-1',
          parentMessageId: 202,
          sequence: 1,
          timestamp: '2026-04-19T00:00:01.000Z',
          event: {
            kind: 'result',
            status: 'completed',
            subagent: {
              id: 501,
              sortOrder: 0,
              name: 'worker-a',
              goal: 'task-1',
              status: 'completed',
              summary: 'done',
              error: null,
            },
          },
        });
        params.subagents.onDone({
          batchId: 'batch-1',
          sessionId: 'session-1',
          parentMessageId: 202,
          results: [
            {
              id: 501,
              sortOrder: 0,
              name: 'worker-a',
              goal: 'task-1',
              status: 'completed',
              summary: 'done',
              error: null,
            },
            {
              id: 502,
              sortOrder: 1,
              name: 'worker-b',
              goal: 'task-2',
              status: 'failed',
              summary: '',
              error: 'boom',
            },
          ],
        });

        return {
          finalContent: 'final answer',
          toolCallCount: 1,
          metadata: {
            existing: true,
          },
        };
      },
    });

    expect(sessionService.updateMessage).toHaveBeenCalledWith(202, {
      content: 'final answer',
      thinking: undefined,
      metadata: expect.objectContaining({
        existing: true,
        inFlight: false,
        subagentBatches: [
          expect.objectContaining({
            batchId: 'batch-1',
            sessionId: 'session-1',
            parentMessageId: 202,
            mode: 'parallel',
            tasks: [
              { goal: 'task-1', agentName: 'worker-a' },
              { goal: 'task-2', agentName: 'worker-b' },
            ],
            status: 'completed',
            results: [
              {
                id: 501,
                sortOrder: 0,
                name: 'worker-a',
                goal: 'task-1',
                status: 'completed',
                summary: 'done',
                error: null,
              },
              {
                id: 502,
                sortOrder: 1,
                name: 'worker-b',
                goal: 'task-2',
                status: 'failed',
                summary: '',
                error: 'boom',
              },
            ],
            metadata: expect.objectContaining({
              events: expect.arrayContaining([
                expect.objectContaining({
                  sequence: 1,
                  timestamp: '2026-04-19T00:00:01.000Z',
                  event: expect.objectContaining({
                    kind: 'result',
                    status: 'completed',
                    subagent: expect.objectContaining({
                      id: 501,
                      summary: 'done',
                    }),
                  }),
                }),
                expect.objectContaining({
                  event: expect.objectContaining({
                    kind: 'status',
                    status: 'completed',
                  }),
                }),
              ]),
              latestEvent: expect.objectContaining({
                event: expect.objectContaining({
                  kind: 'status',
                  status: 'completed',
                }),
              }),
              finalResults: [
                expect.objectContaining({
                  id: 501,
                  summary: 'done',
                }),
                expect.objectContaining({
                  id: 502,
                  error: 'boom',
                }),
              ],
            }),
          }),
        ],
      }),
    });
    expect(result).toEqual({
      sessionId: 'session-1',
      finalContent: 'final answer',
      content: 'final answer',
      usage: undefined,
      toolCallCount: 1,
      assistantMessageId: 202,
      subagentBatches: [
        expect.objectContaining({
          batchId: 'batch-1',
          sessionId: 'session-1',
          parentMessageId: 202,
          mode: 'parallel',
          tasks: [
            { goal: 'task-1', agentName: 'worker-a' },
            { goal: 'task-2', agentName: 'worker-b' },
          ],
          status: 'completed',
          results: [
            {
              id: 501,
              sortOrder: 0,
              name: 'worker-a',
              goal: 'task-1',
              status: 'completed',
              summary: 'done',
              error: null,
            },
            {
              id: 502,
              sortOrder: 1,
              name: 'worker-b',
              goal: 'task-2',
              status: 'failed',
              summary: '',
              error: 'boom',
            },
          ],
          metadata: expect.objectContaining({
            events: expect.any(Array),
            latestEvent: expect.any(Object),
            finalResults: expect.any(Array),
          }),
        }),
      ],
    });
  });

  it('marks the assistant placeholder as failed before rethrowing the runner error', async () => {
    const runnerError = new Error('runner failed');
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([{ id: 101, role: 'user', content: 'hello' }]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'hello' })
        .mockResolvedValueOnce({
          id: 202,
          sessionId: 'session-1',
          role: 'assistant',
          content: '',
          metadata: { inFlight: true },
        }),
      updateMessage: jest.fn().mockResolvedValue(undefined),
    } as unknown as NetaClawSessionService;

    const service = new NetaClawChatOrchestratorService();
    service.sessionService = sessionService;

    await expect(
      service.executeChat({
        sessionId: 'session-1',
        requestedAgentName: 'default',
        message: 'hello',
        runner: jest.fn().mockRejectedValue(runnerError),
      })
    ).rejects.toThrow('runner failed');

    expect(sessionService.updateMessage).toHaveBeenCalledWith(202, {
      content: 'runner failed',
      metadata: {
        inFlight: false,
        error: {
          message: 'runner failed',
          name: 'Error',
        },
      },
    });
  });

  it('rethrows the original runner error when placeholder cleanup also fails', async () => {
    const runnerError = new Error('runner failed');
    const cleanupError = new Error('cleanup failed');
    const sessionService = {
      getOrCreateSession: jest.fn().mockResolvedValue('session-1'),
      loadHistoryWithId: jest.fn().mockResolvedValue([{ id: 101, role: 'user', content: 'hello' }]),
      saveMessage: jest
        .fn()
        .mockResolvedValueOnce({ id: 101, sessionId: 'session-1', role: 'user', content: 'hello' })
        .mockResolvedValueOnce({
          id: 202,
          sessionId: 'session-1',
          role: 'assistant',
          content: '',
          metadata: { inFlight: true },
        }),
      updateMessage: jest.fn().mockRejectedValue(cleanupError),
    } as unknown as NetaClawSessionService;

    const logger = {
      error: jest.fn(),
    } as any;

    const service = new NetaClawChatOrchestratorService() as NetaClawChatOrchestratorService & { logger?: { error: jest.Mock } };
    service.sessionService = sessionService;
    service.logger = logger;

    await expect(
      service.executeChat({
        sessionId: 'session-1',
        requestedAgentName: 'default',
        message: 'hello',
        runner: jest.fn().mockRejectedValue(runnerError),
      })
    ).rejects.toThrow('runner failed');

    expect(logger.error).toHaveBeenCalledWith(
      '[NetaClaw] failed to persist assistant placeholder cleanup: %s',
      'cleanup failed'
    );
  });
});

describe('NetaClawSessionService message persistence', () => {
  let service: NetaClawSessionService;

  beforeEach(() => {
    service = new NetaClawSessionService();
  });

  it('returns the saved message entity and updates it by id', async () => {
    const savedEntity = {
      id: 88,
      sessionId: 'session-1',
      role: 'assistant',
      content: '',
      metadata: { inFlight: true },
    };

    service.messageRepo = {
      save: jest.fn().mockResolvedValue(savedEntity),
      update: jest.fn().mockResolvedValue(undefined),
    } as any;

    const result = await service.saveMessage('session-1', {
      role: 'assistant',
      content: '',
      metadata: { inFlight: true },
    });

    expect(service.messageRepo.save).toHaveBeenCalledWith({
      sessionId: 'session-1',
      role: 'assistant',
      content: '',
      thinking: undefined,
      toolCalls: undefined,
      toolCallId: undefined,
      skillName: undefined,
      metadata: { inFlight: true },
    });
    expect(result).toBe(savedEntity);

    await service.updateMessage(88, {
      content: 'final answer',
      metadata: { inFlight: false, source: 'agent' },
    });

    expect(service.messageRepo.update).toHaveBeenCalledWith(
      { id: 88 },
      {
        content: 'final answer',
        metadata: { inFlight: false, source: 'agent' },
      }
    );
  });
});