openspec/lightweight-tasks/task1-2-5.md

Add tests for runStreamJsonSession and enhance session handling

- Implement tests for runStreamJsonSession to validate user prompts and message handling.
- Improve session termination logic to ensure all active runs are awaited.
- Log user prompts with additional metadata for better tracking.

packages/cli/src/streamJson/session.test.ts

@@ -4,39 +4,238 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { PassThrough } from 'node:stream';
+import { PassThrough, Readable } from 'node:stream';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import type { Config } from '@qwen-code/qwen-code-core';
+import type { LoadedSettings } from '../config/settings.js';
 import { runStreamJsonSession } from './session.js';
 import { StreamJsonController } from './controller.js';
 import { StreamJsonWriter } from './writer.js';
-import type { LoadedSettings } from '../config/settings.js';
+
+const runNonInteractiveMock = vi.fn();
+const logUserPromptMock = vi.fn();
 
 vi.mock('../nonInteractiveCli.js', () => ({
-  runNonInteractive: vi.fn().mockResolvedValue(undefined),
+  runNonInteractive: (...args: unknown[]) => runNonInteractiveMock(...args),
 }));
 
-function createConfig(): Config {
+vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('@qwen-code/qwen-code-core')>();
   return {
+    ...actual,
+    logUserPrompt: (...args: unknown[]) => logUserPromptMock(...args),
+  };
+});
+
+interface ConfigOverrides {
+  getIncludePartialMessages?: () => boolean;
+  getSessionId?: () => string;
+  getModel?: () => string;
+  getContentGeneratorConfig?: () => { authType?: string };
+  [key: string]: unknown;
+}
+
+function createConfig(overrides: ConfigOverrides = {}): Config {
+  const base = {
     getIncludePartialMessages: () => false,
     getSessionId: () => 'session-test',
     getModel: () => 'model-test',
-  } as unknown as Config;
+    getContentGeneratorConfig: () => ({ authType: 'test-auth' }),
+    getOutputFormat: () => 'stream-json',
+  };
+  return { ...base, ...overrides } as unknown as Config;
+}
+
+function createSettings(): LoadedSettings {
+  return {
+    merged: {
+      security: { auth: {} },
+    },
+  } as unknown as LoadedSettings;
+}
+
+function createWriter() {
+  return {
+    emitResult: vi.fn(),
+    writeEnvelope: vi.fn(),
+    emitSystemMessage: vi.fn(),
+  } as unknown as StreamJsonWriter;
 }
 
 describe('runStreamJsonSession', () => {
   let settings: LoadedSettings;
 
   beforeEach(() => {
-    vi.spyOn(process.stdout, 'write').mockImplementation(() => true);
-    settings = {} as LoadedSettings;
+    settings = createSettings();
+    runNonInteractiveMock.mockReset();
+    logUserPromptMock.mockReset();
   });
 
   afterEach(() => {
     vi.restoreAllMocks();
   });
 
-  it('delegates incoming control requests to the controller', async () => {
+  it('runs initial prompt before reading stream and logs it', async () => {
+    const config = createConfig();
+    const writer = createWriter();
+    const stream = Readable.from([]);
+    runNonInteractiveMock.mockResolvedValueOnce(undefined);
+
+    await runStreamJsonSession(config, settings, 'Hello world', {
+      input: stream,
+      writer,
+    });
+
+    expect(runNonInteractiveMock).toHaveBeenCalledTimes(1);
+    const call = runNonInteractiveMock.mock.calls[0];
+    expect(call[0]).toBe(config);
+    expect(call[1]).toBe(settings);
+    expect(call[2]).toBe('Hello world');
+    expect(typeof call[3]).toBe('string');
+    expect(call[4]).toEqual(
+      expect.objectContaining({
+        streamJson: expect.objectContaining({ writer }),
+        abortController: expect.any(AbortController),
+      }),
+    );
+    expect(logUserPromptMock).toHaveBeenCalledTimes(1);
+    const loggedPrompt = logUserPromptMock.mock.calls[0][1] as
+      | Record<string, unknown>
+      | undefined;
+    expect(loggedPrompt).toMatchObject({
+      prompt: 'Hello world',
+      prompt_length: 11,
+    });
+    expect(loggedPrompt?.['prompt_id']).toBe(call[3]);
+  });
+
+  it('handles user envelope when no initial prompt is provided', async () => {
+    const config = createConfig();
+    const writer = createWriter();
+    const envelope = {
+      type: 'user' as const,
+      message: {
+        content: '   Stream mode ready   ',
+      },
+    };
+    const stream = Readable.from([`${JSON.stringify(envelope)}\n`]);
+    runNonInteractiveMock.mockResolvedValueOnce(undefined);
+
+    await runStreamJsonSession(config, settings, undefined, {
+      input: stream,
+      writer,
+    });
+
+    expect(runNonInteractiveMock).toHaveBeenCalledTimes(1);
+    const call = runNonInteractiveMock.mock.calls[0];
+    expect(call[2]).toBe('Stream mode ready');
+    expect(call[4]).toEqual(
+      expect.objectContaining({
+        userEnvelope: envelope,
+        streamJson: expect.objectContaining({ writer }),
+        abortController: expect.any(AbortController),
+      }),
+    );
+  });
+
+  it('processes multiple user messages sequentially', async () => {
+    const config = createConfig();
+    const writer = createWriter();
+    const lines = [
+      JSON.stringify({
+        type: 'user',
+        message: { content: 'first request' },
+      }),
+      JSON.stringify({
+        type: 'user',
+        message: { content: 'second request' },
+      }),
+    ].map((line) => `${line}\n`);
+    const stream = Readable.from(lines);
+    runNonInteractiveMock.mockResolvedValue(undefined);
+
+    await runStreamJsonSession(config, settings, undefined, {
+      input: stream,
+      writer,
+    });
+
+    expect(runNonInteractiveMock).toHaveBeenCalledTimes(2);
+    expect(runNonInteractiveMock.mock.calls[0][2]).toBe('first request');
+    expect(runNonInteractiveMock.mock.calls[1][2]).toBe('second request');
+  });
+
+  it('emits stream_event when partial messages are enabled', async () => {
+    const config = createConfig({
+      getIncludePartialMessages: () => true,
+      getSessionId: () => 'partial-session',
+      getModel: () => 'partial-model',
+    });
+    const stream = Readable.from([
+      `${JSON.stringify({
+        type: 'user',
+        message: { content: 'show partial' },
+      })}\n`,
+    ]);
+    const writeSpy = vi
+      .spyOn(process.stdout, 'write')
+      .mockImplementation(() => true);
+
+    runNonInteractiveMock.mockImplementationOnce(
+      async (
+        _config,
+        _settings,
+        _prompt,
+        _promptId,
+        options?: {
+          streamJson?: { writer?: StreamJsonWriter };
+        },
+      ) => {
+        const builder = options?.streamJson?.writer?.createAssistantBuilder();
+        builder?.appendText('partial');
+        builder?.finalize();
+      },
+    );
+
+    await runStreamJsonSession(config, settings, undefined, {
+      input: stream,
+    });
+
+    const outputs = writeSpy.mock.calls
+      .map(([chunk]) => chunk as string)
+      .join('')
+      .split('\n')
+      .map((line) => line.trim())
+      .filter((line) => line.length > 0)
+      .map((line) => JSON.parse(line));
+
+    expect(outputs.some((envelope) => envelope.type === 'stream_event')).toBe(
+      true,
+    );
+    writeSpy.mockRestore();
+  });
+
+  it('emits error result when JSON parsing fails', async () => {
+    const config = createConfig();
+    const writer = createWriter();
+    const stream = Readable.from(['{invalid json\n']);
+
+    await runStreamJsonSession(config, settings, undefined, {
+      input: stream,
+      writer,
+    });
+
+    expect(writer.emitResult).toHaveBeenCalledWith(
+      expect.objectContaining({
+        isError: true,
+      }),
+    );
+    expect(runNonInteractiveMock).not.toHaveBeenCalled();
+  });
+
+  it('delegates control requests to the controller', async () => {
+    const config = createConfig();
+    const writer = new StreamJsonWriter(config, false);
     const controllerPrototype = StreamJsonController.prototype as unknown as {
       handleIncomingControlRequest: (...args: unknown[]) => unknown;
     };
@@ -46,8 +245,6 @@ describe('runStreamJsonSession', () => {
     );
 
     const inputStream = new PassThrough();
-    const config = createConfig();
-
     const controlRequest = {
       type: 'control_request',
       request_id: 'req-1',
@@ -58,7 +255,7 @@ describe('runStreamJsonSession', () => {
 
     await runStreamJsonSession(config, settings, undefined, {
       input: inputStream,
-      writer: new StreamJsonWriter(config, false),
+      writer,
     });
 
     expect(handleSpy).toHaveBeenCalledTimes(1);