1
10import { describe, expect, it } from "vitest";
11
12import { MockProvider, parseWritingBlocks, writingTransform } from "./index.js";
13import type { ChatRequest } from "../types.js";
14
15function writingReq(opts: { action: string; blocks: Array<{ id: string; text: string }>; language?: string } & Record<string, unknown>): ChatRequest {
16 const userText =
17 "Do the writing task.\n" +
18 opts.blocks
19 .map((b) => `[SELECTED ${b.id}]\n${b.text}\n[/SELECTED ${b.id}]`)
20 .join("\n");
21 const metadata: Record<string, string> = { mode: "writing", action: opts.action };
22 if (opts.language) metadata.language = opts.language;
23 return {
24 model: "mock",
25 messages: [
26 { role: "system", content: "You generate block edits." },
27 { role: "user", content: userText },
28 ],
29 tools: [
30 {
31 name: "generateBlocks",
32 description: "Emit the block edits.",
33 parameters: {
34 type: "object",
35 properties: {
36 blocks: {
37 type: "array",
38 items: {
39 type: "object",
40 properties: { blockId: { type: "string" }, newText: { type: "string" } },
41 required: ["blockId", "newText"],
42 },
43 },
44 },
45 required: ["blocks"],
46 },
47 },
48 ],
49 toolChoice: { name: "generateBlocks" },
50 metadata: { ...metadata, ...(opts as Record<string, unknown>) } as unknown as Record<string, string>,
51 };
52}
53
54async function collect(req: ChatRequest): Promise<Record<string, unknown>[]> {
55 const p = new MockProvider("openai", 8);
56 const out: Record<string, unknown>[] = [];
57 for await (const chunk of p.chat(req as never)) out.push(chunk as Record<string, unknown>);
58 return out;
59}
60
61describe("MockProvider writing mode (TASK-0123)", () => {
62 it("summarize: one tool call, per-block 'Summarized:' of the first 80 chars", async () => {
63 const longText = "x".repeat(120);
64 const chunks = await collect(
65 writingReq({ action: "summarize", blocks: [{ id: "b1", text: longText }, { id: "b2", text: "short" }] }),
66 );
67 const call = chunks.find((c) => c.toolCallDelta)!.toolCallDelta! as {
68 id: string;
69 name: string;
70 arguments: { blocks: Array<{ blockId: string; newText: string }> };
71 };
72 expect(call.name).toBe("generateBlocks");
73 expect(call.arguments.blocks).toEqual([
74 { blockId: "b1", newText: `Summarized: ${"x".repeat(80)}` },
75 { blockId: "b2", newText: "Summarized: short" },
76 ]);
77 const done = chunks.at(-1) as { usage: { in: number; out: number }; finishReason: string; done: boolean };
78 expect(done.done).toBe(true);
79 expect(done.finishReason).toBe("tool_calls");
80 expect(done.usage.in).toBeGreaterThan(0);
81 expect(done.usage.out).toBeGreaterThan(0);
82 });
83
84 it("rewrite / improve / translate transforms are deterministic", async () => {
85 const rewrite = await collect(writingReq({ action: "rewrite", blocks: [{ id: "b1", text: "hello world" }] }));
86 expect(
87 (rewrite.find((c) => c.toolCallDelta)!.toolCallDelta! as { arguments: { blocks: Array<{ blockId: string; newText: string }> } }).arguments.blocks,
88 ).toEqual([{ blockId: "b1", newText: "Rewritten: hello world" }]);
89
90 const improve = await collect(writingReq({ action: "improve", blocks: [{ id: "b1", text: "meh sentence" }] }));
91 expect(
92 (improve.find((c) => c.toolCallDelta)!.toolCallDelta! as { arguments: { blocks: Array<{ blockId: string; newText: string }> } }).arguments.blocks,
93 ).toEqual([{ blockId: "b1", newText: "Improved: meh sentence" }]);
94
95 const translate = await collect(writingReq({ action: "translate", blocks: [{ id: "b1", text: "bonjour" }], language: "fr" }));
96 expect(
97 (translate.find((c) => c.toolCallDelta)!.toolCallDelta! as { arguments: { blocks: Array<{ blockId: string; newText: string }> } }).arguments.blocks,
98 ).toEqual([{ blockId: "b1", newText: "Translated to fr: bonjour" }]);
99 });
100
101 it("preserves prompt order across multiple selected blocks (incl. interleaved context markers)", async () => {
102 const req = writingReq({ action: "rewrite", blocks: [{ id: "b3", text: "c" }, { id: "b1", text: "a" }, { id: "b2", text: "b" }] });
103
104 const lastUser = req.messages[req.messages.length - 1]!;
105 lastUser.content = `[CONTEXT x0]\nctx\n[/CONTEXT x0]\n[SELECTED b3]\nc\n[/SELECTED b3]\n[CONTEXT x1]\nctx2\n[/CONTEXT x1]\n[SELECTED b1]\na\n[/SELECTED b1]\n[SELECTED b2]\nb\n[/SELECTED b2]`;
106 const chunks = await collect(req);
107 const call = chunks.find((c) => c.toolCallDelta)!.toolCallDelta! as { arguments: { blocks: Array<{ blockId: string }> } };
108 expect(call.arguments.blocks.map((b) => b.blockId)).toEqual(["b3", "b1", "b2"]);
109 });
110
111 it("metadata.forceUnknownId appends an unknown blockId (validation seam)", async () => {
112 const chunks = await collect(
113 writingReq({ action: "rewrite", blocks: [{ id: "b1", text: "t" }], forceUnknownId: "not_in_prompt" }),
114 );
115 const call = chunks.find((c) => c.toolCallDelta)!.toolCallDelta! as { arguments: { blocks: Array<{ blockId: string }> } };
116 expect(call.arguments.blocks.map((b) => b.blockId)).toEqual(["b1", "not_in_prompt"]);
117 });
118
119 it("armed failAfterDeltas dies mid-stream (no tool call, no done frame)", async () => {
120 const p = new MockProvider("openai", 8);
121 p.failAfterDeltas = 2;
122 p.failStatus = 503;
123 const chunks: unknown[] = [];
124 let threw = false;
125 try {
126 for await (const c of p.chat(writingReq({ action: "rewrite", blocks: [{ id: "b1", text: "t" }] }) as never)) chunks.push(c);
127 } catch {
128 threw = true;
129 }
130 expect(threw).toBe(true);
131 expect(chunks.some((c) => (c as { toolCallDelta?: unknown }).toolCallDelta)).toBe(false);
132 expect(chunks.some((c) => (c as { done?: boolean }).done === true)).toBe(false);
133
134 const ok = await collect(
135 writingReq({ action: "rewrite", blocks: [{ id: "b1", text: "again" }] }),
136 );
137 expect(ok.some((c) => (c as { toolCallDelta?: unknown }).toolCallDelta)).toBe(true);
138 });
139
140 it("deltaDelayMs paces the call (deterministic latency floor)", async () => {
141 const p = new MockProvider("openai", 8);
142 p.deltaDelayMs = 40;
143 const started = Date.now();
144 for await (const c of p.chat(writingReq({ action: "rewrite", blocks: [{ id: "b1", text: "t" }] }) as never)) {
145 if ((c as { toolCallDelta?: unknown }).toolCallDelta) break;
146 }
147 expect(Date.now() - started).toBeGreaterThanOrEqual(30);
148 });
149});
150
151describe("parseWritingBlocks / writingTransform (pure helpers)", () => {
152 it("parses only SELECTED markers, in order, ignoring CONTEXT", () => {
153 const msgs = [
154 {
155 role: "user",
156 content: "[CONTEXT a]\nctx\n[/CONTEXT a]\n[SELECTED s1]\ntext one\n[/SELECTED s1]\n[SELECTED s2]\ntwo\n[/SELECTED s2]",
157 },
158 ];
159 expect(parseWritingBlocks(msgs)).toEqual([
160 { id: "s1", text: "text one" },
161 { id: "s2", text: "two" },
162 ]);
163 });
164
165 it("returns [] with no user message / no markers", () => {
166 expect(parseWritingBlocks([])).toEqual([]);
167 expect(parseWritingBlocks([{ role: "system", content: "[SELECTED x]\nno\n[/SELECTED x]" }])).toEqual([]);
168 expect(parseWritingBlocks([{ role: "user", content: "no markers here" }])).toEqual([]);
169 });
170
171 it("transforms are stable per action", () => {
172 expect(writingTransform("summarize", "a".repeat(100))).toBe(`Summarized: ${"a".repeat(80)}`);
173 expect(writingTransform("improve", " x ")).toBe("Improved: x");
174 expect(writingTransform("translate", "y")).toBe("Translated to the target language: y");
175 expect(writingTransform("bogus", "z")).toBe("Rewritten: z");
176 });
177});
178
