fix(llm): keep system prompts on premium chat, reject incomplete streams, stop double-charging quota

apps/desktop/tests/main/services/local-llm-chat-contract-r1-6.test.ts

@@ -0,0 +1,100 @@
+// LocalLLMService.chatStream 이 @d3ro/core/llm-chat 계약을 따르는지 (red-team r1-6)
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { ErrorCode } from '@d3ro/core/errors'
+import { initInMemoryConfig, resetInMemoryConfig } from '../../../src/main/services/ConfigService'
+import {
+  getLocalLLMService,
+  resetLocalLLMServiceForTests,
+} from '../../../src/main/services/LocalLLMService'
+
+const encoder = new TextEncoder()
+
+function ndjsonResponse(parts: string[]): Response {
+  return new Response(new ReadableStream<Uint8Array>({
+    start(controller) {
+      for (const part of parts) controller.enqueue(encoder.encode(part))
+      controller.close()
+    },
+  }), { status: 200 })
+}
+
+async function collect(gen: AsyncGenerator<string, string>): Promise<{ tokens: string[]; result: string }> {
+  const tokens: string[] = []
+  for (let next = await gen.next(); ; next = await gen.next()) {
+    if (next.done) return { tokens, result: next.value }
+    tokens.push(next.value)
+  }
+}
+
+describe('LocalLLMService chat contract', () => {
+  beforeEach(() => {
+    initInMemoryConfig()
+    resetLocalLLMServiceForTests()
+    ;(getLocalLLMService() as unknown as { _available: boolean })._available = true
+  })
+
+  afterEach(() => {
+    resetLocalLLMServiceForTests()
+    resetInMemoryConfig()
+    vi.unstubAllGlobals()
+  })
+
+  it('sends the system prompt as a single leading system message', async () => {
+    let sent: Array<{ role: string; content: string }> = []
+    vi.stubGlobal('fetch', vi.fn(async (_input: RequestInfo | URL, init?: RequestInit) => {
+      sent = (JSON.parse(String(init?.body)) as { messages: typeof sent }).messages
+      return ndjsonResponse([
+        `${JSON.stringify({ message: { content: 'hi' }, done: false })}\n`,
+        JSON.stringify({ message: { content: '' }, done: true }),
+      ])
+    }))
+
+    const { tokens, result } = await collect(getLocalLLMService().chatStream([
+      { role: 'system', content: 'persona' },
+      { role: 'user', content: 'hello' },
+    ]))
+
+    expect(tokens).toEqual(['hi'])
+    expect(result).toBe('hi')
+    expect(sent).toEqual([
+      { role: 'system', content: 'persona' },
+      { role: 'user', content: 'hello' },
+    ])
+  })
+
+  it('rejects a stream that ends without a done frame', async () => {
+    vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse([
+      `${JSON.stringify({ message: { content: 'partial' }, done: false })}\n`,
+    ])))
+
+    await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
+      .rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed })
+  })
+
+  it('rejects malformed NDJSON frames, including non-object JSON', async () => {
+    vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse(['{broken\n'])))
+    await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
+      .rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed, message: 'Ollama returned malformed NDJSON' })
+
+    vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse(['null\n'])))
+    await expect(collect(getLocalLLMService().chatStream([{ role: 'user', content: 'q' }])))
+      .rejects.toMatchObject({ code: ErrorCode.LLMProcessingFailed, message: 'Ollama returned malformed NDJSON' })
+  })
+
+  it('streamGenerate still completes on a trailing done frame without newline', async () => {
+    vi.stubGlobal('fetch', vi.fn(async () => ndjsonResponse([
+      `${JSON.stringify({ model: 'm', response: 'a', done: false })}\n`,
+      JSON.stringify({ model: 'm', response: 'b', done: true, eval_count: 2 }),
+    ])))
+
+    const gen = getLocalLLMService().streamGenerate('p')
+    const tokens: string[] = []
+    let next = await gen.next()
+    while (!next.done) {
+      tokens.push(next.value)
+      next = await gen.next()
+    }
+    expect(tokens).toEqual(['a', 'b'])
+    expect(next.value).toMatchObject({ text: 'ab', completionTokens: 2 })
+  })
+})