Merge main: wave-4 backend (activity abort endpoint, lastStep capture)

2026-09-11 03:33:15 +03:00
parent 99ba029a3c 090a39b77b
commit 3b1d0ed18d
21 changed files with 1021 additions and 65 deletions
@@ -48,7 +48,7 @@ import { runDemoBuilder } from "../src/agent/subagents/demo-builder.ts";
 import { headlessOverride, type DemoRunResult, type HeadlessApi } from "../src/agent/subagents/headless.ts";
 import { createJobRunner, demoModelOverride, visualModelOverride } from "../src/agent/subagents/runner.ts";
 import { runVisualMaker } from "../src/agent/subagents/visual-maker.ts";
-import { buildSubagentSteer, isSubagentFailure, type SubagentSteerMessage } from "../src/core/contracts.ts";
+import { buildAbortSteer, buildSubagentSteer, isSubagentFailure, type SubagentSteerMessage } from "../src/core/contracts.ts";
 import type { JobRecord } from "../src/core/contracts.ts";
 import { createTeacherTools } from "../src/agent/tools.ts";
 import { chatStore } from "../src/store/chats.ts";
@@ -341,6 +341,18 @@ describe("steer card rendering (the delivered artifact)", () => {
 		assert.ok(!html.includes("Looking at the render"), "model-facing notes never render");
 	});

+	it("renders the abort chip for an aborted job's steer (2026-09 dock semantics)", () => {
+		const steer = buildAbortSteer({ id: "jc", kind: "research", toolCallId: "cc" });
+		assert.equal(steer.metadata.subagentResult.aborted, true);
+		assert.equal(steer.parts.length, 1, "no artifact part on an abort");
+		const html = renderToStaticMarkup(<SubagentSteerCard message={steer as never} />);
+		assert.match(html, /job-aborted-badge/);
+		assert.match(html, /research/, "the job's kind, lower-cased like the complete-badge idiom");
+		assert.match(html, /aborted by the learner/);
+		// the teacher-facing protocol text never renders for the learner
+		assert.ok(!html.includes("no brief is coming"), "no protocol glue in the chip");
+	});
+
 	it("renders a quiet Research-complete one-liner whose collapsed <details> carries the full brief", () => {
 		const steer = buildSubagentSteer(
 			{ id: "jr", kind: "research", toolCallId: "cr" },
@@ -34,7 +34,7 @@ import type {
 	SessionFindings,
 } from "../src/core/contracts.ts";
 import type { GraphStore } from "../src/core/ports.ts";
-import { createExtractionPipeline, extractionStatuses, transcriptView } from "../src/agent/subagents/extraction.ts";
+import { abortExtraction, createExtractionPipeline, extractionStatuses, transcriptView } from "../src/agent/subagents/extraction.ts";
 import { chatStore } from "../src/store/chats.ts";
 import { PATHS } from "../config.ts";

@@ -197,6 +197,24 @@ function submitModel(name: string, input: unknown, onCall?: () => void): MockLan
 	});
 }

+/** A model that hangs until its abort signal fires, then rejects —
+ *  mimicking a cancelled provider request (the abort path). Like a real
+ *  provider, an ALREADY-aborted signal rejects immediately (the abort
+ *  can win the race before doGenerate attaches its listener). */
+function hangingModel(): MockLanguageModelV4 {
+	const abortReject = (reject: (reason: Error) => void) =>
+		reject(Object.assign(new Error("The operation was aborted"), { name: "AbortError" }));
+	return new MockLanguageModelV4({
+		doGenerate: async (options): Promise<never> => {
+			await new Promise<never>((_, reject) => {
+				if (options.abortSignal?.aborted) abortReject(reject);
+				else options.abortSignal?.addEventListener("abort", () => abortReject(reject));
+			});
+			throw new Error("unreachable");
+		},
+	});
+}
+
 /** A model that never calls a tool (the failure path). */
 function silentModel(): MockLanguageModelV4 {
 	return new MockLanguageModelV4({
@@ -378,6 +396,97 @@ describe("the extraction pipeline (extractor → auditor → merger)", () => {
 		assert.equal(extractionStatuses().length, 0, "the registry clears when the run ends");
 	});

+	it("captures per-stage usage live, persists the run's usage into the chat record", async () => {
+		const { chatId } = fixtureChat();
+		const ledgers = memoryLedgers();
+		const existingId = "n_existing";
+		ledgers.appendNode({
+			id: existingId,
+			name: "Derivative",
+			statement: "Derivatives measure rates of change",
+			aliases: [],
+			createdAt: "2026-08-01T00:00:00.000Z",
+			provenance: [],
+		});
+		const usageSeen: (number | undefined)[] = [];
+		const verdict: AuditResult = {
+			propositions: [{ ...FINDINGS.propositions[0]!, linkTo: existingId }],
+			edges: [],
+			attribution: { quiz: {} },
+			dropped: [{ ref: "p2", reason: "evidence not found" }],
+		};
+		const pipeline = createExtractionPipeline({
+			chatStore,
+			graphStore: ledgers,
+			// During the auditor's generate, the registry must already carry
+			// the extractor's fully accumulated live usage (2 steps × 10/5).
+			extractor: submitModel("submit_findings", FINDINGS),
+			auditor: submitModel("submit_review", verdict, () =>
+				usageSeen.push(extractionStatuses().find((r) => r.chatId === chatId)?.usage.extractor?.steps.at(-1)?.inputTokens),
+			),
+		});
+		await pipeline.runExtraction(chatId);
+
+		// Live: mid-audit the extractor's usage was fully accumulated
+		// (the mock bills 10 input / 5 output per step, flat per the SDK).
+		assert.deepEqual(usageSeen, [10]);
+
+		// Persisted: the run appended ONE usage entry to the chat's
+		// extractions log — per-stage model + steps, with merge stats.
+		const record = chatStore.getChat(chatId)!;
+		assert.equal(record.extractions?.length, 1);
+		const run = record.extractions![0]!;
+		assert.equal(run.stats?.nodes, 0); // linked, not created
+		assert.equal(run.stats?.events, 3);
+		assert.equal(run.stages.extractor.model, "mock-model-id"); // the mock's model id
+		assert.equal(run.stages.extractor.steps.length, 2); // tool-call step + text step
+		assert.deepEqual(run.stages.extractor.steps[0], {
+			inputTokens: 10,
+			outputTokens: 5,
+			cacheReadTokens: 0,
+			cacheWriteTokens: 0,
+			reasoningTokens: 0,
+		});
+		assert.equal(run.stages.auditor.steps.length, 2);
+		assert.ok(!Number.isNaN(Date.parse(run.at)));
+	});
+
+	it("aborts mid-run: cancels the model call, persists partial usage, checkpoints nothing", async () => {
+		const { chatId } = fixtureChat();
+		const ledgers = memoryLedgers();
+
+		const pipeline = createExtractionPipeline({
+			chatStore,
+			graphStore: ledgers,
+			extractor: hangingModel(), // hangs until aborted — no step finishes
+			auditor: submitModel("submit_review", { propositions: [], edges: [], attribution: { quiz: {} }, dropped: [] }),
+		});
+		const runPromise = pipeline.runExtraction(chatId);
+
+		// The registry entry exists synchronously after the call (setPhase
+		// runs before the first await); abort is TRUE cancellation.
+		assert.equal(abortExtraction(chatId), true);
+		const stats = await runPromise;
+
+		// A deliberate stop, not a failure: zero emissions, no throw.
+		assert.deepEqual(stats, { nodes: 0, edges: 0, events: 0 });
+		assert.equal(extractionStatuses().length, 0, "the registry clears");
+		assert.equal(ledgers.readNodes().length, 0, "nothing merged");
+		assert.equal(ledgers.readCheckpoints().length, 0, "nothing checkpointed — the next trigger re-runs");
+
+		// Partial usage still persisted (money was spent): stats absent,
+		// no completed steps in either stage.
+		const record = chatStore.getChat(chatId)!;
+		assert.equal(record.extractions?.length, 1);
+		const run = record.extractions![0]!;
+		assert.equal(run.stats, undefined);
+		assert.equal(run.stages.extractor.steps.length, 0);
+		assert.equal(run.stages.auditor.steps.length, 0);
+
+		// Aborting a non-running chat is an idempotent no-op.
+		assert.equal(abortExtraction(chatId), false);
+	});
+
 	it("is single-flight per chat — a concurrent second run is a no-op", async () => {
 		const { chatId } = fixtureChat();
 		const ledgers = memoryLedgers();
@@ -31,18 +31,23 @@ import type { UIMessage } from "ai";
 import { MockLanguageModelV4, simulateReadableStream } from "ai/test";
 import type { LanguageModelV4StreamPart } from "@ai-sdk/provider";
 import {
+	accumulateJobUsage,
+	buildAbortSteer,
 	buildSubagentSteer,
 	isSubagentFailure,
 	steerDelivered,
 	type DemoResult,
+	type JobUsage,
 	type ResearchResult,
+	type SteerMessage,
 	type SubagentSteerMessage,
 	type VisualResult,
 } from "../src/core/contracts.ts";
 import { composeInstructions, createTeacherAgent } from "../src/agent/teacher-agent.ts";
 import { createTeacherTools } from "../src/agent/tools.ts";
 import { handleChatRequest } from "../src/server/chat-handler.ts";
-import { claimSteers, deliverJobResult, setTurnRunner } from "../src/server/jobs.ts";
+import { abortJob, claimSteers, deliverJobResult, setTurnRunner } from "../src/server/jobs.ts";
+import { beginTurn, endTurn } from "../src/server/turn-liveness.ts";
 import { persistTurn, withMessageIds } from "../src/server/turn-service.ts";
 import { chatStore } from "../src/store/chats.ts";
 import { jobStore } from "../src/store/jobs.ts";
@@ -170,6 +175,57 @@ describe("steer helpers (contracts.ts)", () => {
 		assert.deepEqual(data.data.result, research);
 	});

+	it("accumulateJobUsage: sums per model call; fields NO call reported stay undefined (absent is not zero)", () => {
+		let total: JobUsage | undefined = undefined;
+		total = accumulateJobUsage(total, "z-ai/glm-5.3-flash", {
+			inputTokens: 10,
+			outputTokens: 5,
+			inputTokenDetails: { cacheReadTokens: 4 },
+		});
+		assert.deepEqual(total, {
+			model: "z-ai/glm-5.3-flash",
+			inputTokens: 10,
+			outputTokens: 5,
+			cacheReadTokens: 4,
+			cacheWriteTokens: undefined,
+			reasoningTokens: undefined,
+			steps: 1,
+			lastStep: {
+				inputTokens: 10,
+				outputTokens: 5,
+				cacheReadTokens: 4,
+				cacheWriteTokens: undefined,
+				reasoningTokens: undefined,
+			},
+		}, "lastStep is the most recent call's accounting, not cumulative");
+		// a later call that reports nothing new keeps reported fields summing
+		// (unreported contributes 0) and bumps the billing step count.
+		total = accumulateJobUsage(total, "z-ai/glm-5.3-flash", { inputTokens: 20, outputTokens: 0 });
+		assert.equal(total.inputTokens, 30);
+		assert.equal(total.outputTokens, 5);
+		assert.equal(total.cacheReadTokens, 4, "unreported after reported contributes 0");
+		assert.equal(total.cacheWriteTokens, undefined, "never reported stays absent — never zero");
+		assert.equal(total.reasoningTokens, undefined);
+		assert.equal(total.steps, 2, "one step per model call, reported or not");
+		assert.deepEqual(total.lastStep, {
+			inputTokens: 20,
+			outputTokens: 0,
+			cacheReadTokens: undefined,
+			cacheWriteTokens: undefined,
+			reasoningTokens: undefined,
+		}, "lastStep always mirrors the LATEST call, never the sums");
+	});
+
+	it("buildAbortSteer: metadata-aborted, text-only (no artifact), same deterministic id space as the result steer", () => {
+		const steer = buildAbortSteer({ id: "job_c", kind: "demo", toolCallId: "call-c" });
+		assert.equal(steer.id, "steer_job_c", "same id space — steerDelivered treats either as delivered");
+		assert.equal(steer.role, "user");
+		assert.equal(steer.metadata.subagentResult.aborted, true);
+		assert.equal(steer.metadata.subagentResult.jobId, "job_c");
+		assert.equal(steer.parts.length, 1);
+		assert.ok(steer.parts[0]!.type === "text" && /aborted/.test(steer.parts[0].text));
+	});
+
 	it("isSubagentFailure: failure shapes only (empty html/svg, failure brief)", () => {
 		assert.equal(isSubagentFailure("research", research), false);
 		assert.equal(isSubagentFailure("research", { brief: "Research failed: x", sources: [] }), true);
@@ -836,3 +892,192 @@ describe("prompt composition", () => {
 		assert.ok(composeInstructions().length <= 60_000, "composed system prompt over the 60k guard");
 	});
 });
+
+// ─── abort (2026-09 semantics: real kill, cancelled status, teacher told) ───
+
+describe("abort (src/server/jobs.ts abortJob)", () => {
+	/** A hanging model that only settles when its abort signal fires — the
+	 *  real-kill path: the in-flight call itself is killed. */
+	function hangingModel(): MockLanguageModelV4 {
+		return new MockLanguageModelV4({
+			doGenerate: ({ abortSignal }) =>
+				new Promise((_resolve, reject) => {
+					// an ALREADY-aborted signal never fires the event — real
+					// providers reject an aborted fetch at once, the mock must too
+					if (abortSignal?.aborted) {
+						reject(new DOMException("The operation was aborted", "AbortError"));
+						return;
+					}
+					abortSignal?.addEventListener("abort", () => {
+						reject(new DOMException("The operation was aborted", "AbortError"));
+					});
+				}),
+		});
+	}
+
+	/** One step (with usage) then a hang — proves usage-so-far survives the kill. */
+	function oneStepThenHangModel(): MockLanguageModelV4 {
+		let turn = 0;
+		return new MockLanguageModelV4({
+			doGenerate: ({ abortSignal }) => {
+				turn++;
+				if (turn === 1) {
+					return Promise.resolve({
+						content: [
+							{
+								type: "tool-call" as const,
+								toolCallId: "ws1",
+								toolName: "web_search",
+								input: JSON.stringify({ query: "anything" }),
+							},
+						],
+						finishReason: { unified: "tool-calls" as const, raw: "tool-calls" },
+						usage: usage(),
+						warnings: [],
+					});
+				}
+				return new Promise((_resolve, reject) => {
+					if (abortSignal?.aborted) {
+						reject(new DOMException("The operation was aborted", "AbortError"));
+						return;
+					}
+					abortSignal?.addEventListener("abort", () => {
+						reject(new DOMException("The operation was aborted", "AbortError"));
+					});
+				});
+			},
+		});
+	}
+
+	function newChatWithResearchJob(label: string, toolCallId: string): { chatId: string; jobId: string } {
+		const chat = chatStore.createChat(`Test: ${label}`);
+		trackChat(chat.id);
+		chatStore.saveMessages(chat.id, [
+			{ id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+			{
+				id: "a1",
+				role: "assistant",
+				parts: [
+					{ type: "text", text: "Checking." },
+					{ type: "tool-research", toolCallId, state: "output-available", input: RESEARCH_INPUT, output: "Dispatched." },
+				],
+			},
+		] as never);
+		const job = jobStore.createJob({ kind: "research", chatId: chat.id, toolCallId });
+		trackJob(job.id);
+		return { chatId: chat.id, jobId: job.id };
+	}
+
+	it("kills the running loop: error \"Aborted by learner\" immediately, loop unwinds silently, teacher gets the abort steer", async () => {
+		const { jobId } = newChatWithResearchJob("abort running job", "call-abort1");
+		const fired: SteerMessage[] = [];
+		setTurnRunner(async (_cid, messages) => {
+			fired.push(messages[messages.length - 1] as SteerMessage);
+			return 200;
+		});
+		researchModelOverride.current = hangingModel();
+		const runner = createJobRunner({
+			chatStore,
+			jobStore,
+			deliver: async () => {
+				throw new Error("a cancelled job must never deliver a result steer");
+			},
+		});
+		try {
+			const settled = runner.runJob(jobId, RESEARCH_INPUT);
+			await waitFor(() => jobStore.getJob(jobId)!.activity === "researching", 8000, "loop started");
+
+			const outcome = abortJob(jobId);
+			assert.ok(outcome.ok);
+			const killed = jobStore.getJob(jobId)!;
+			assert.equal(killed.status, "error", "flipped synchronously — the poll sees it at once");
+			assert.equal(killed.error, "Aborted by learner", "abort reads as an error with a reason, not a failure");
+			assert.equal(killed.activity, "aborted");
+
+			await settled; // unwinds — no hang
+			assert.equal(jobStore.getJob(jobId)!.status, "error", "the runner never overwrites the flip");
+			assert.equal(jobStore.getJob(jobId)!.error, "Aborted by learner");
+			assert.equal(jobStore.getJob(jobId)!.result, undefined, "no result on an aborted job");
+
+			await waitFor(() => fired.length > 0, 8000, "abort wake fired");
+			const last = fired[0]!;
+			assert.equal(last.metadata.subagentResult.aborted, true);
+			assert.equal(last.id, `steer_${jobId}`);
+			assert.equal(last.parts.length, 1, "text only — no artifact part");
+			// aborted is terminal: the wake completing must not flip it
+			assert.equal(jobStore.getJob(jobId)!.status, "error", "terminal after the wake");
+		} finally {
+			researchModelOverride.current = undefined;
+		}
+	});
+
+	it("the usage folded so far survives the kill (money channel outlives the abort)", async () => {
+		const restoreFetch = stubFetch();
+		try {
+			const { jobId } = newChatWithResearchJob("abort keeps usage", "call-abort2");
+			setTurnRunner(async () => 200);
+			researchModelOverride.current = oneStepThenHangModel();
+			const runner = createJobRunner({ chatStore, jobStore, deliver: async () => {} });
+			try {
+				const settled = runner.runJob(jobId, RESEARCH_INPUT);
+				await waitFor(() => (jobStore.getJob(jobId)!.usage?.steps ?? 0) === 1, 8000, "first step folded");
+				abortJob(jobId);
+				await settled;
+				const job = jobStore.getJob(jobId)!;
+				assert.equal(job.status, "error");
+			assert.equal(job.error, "Aborted by learner");
+				assert.equal(job.usage!.steps, 1, "usage-so-far persisted by the runner's finally");
+				assert.equal(job.usage!.model, "z-ai/glm-5.3-flash", "the config role's OpenRouter id");
+				assert.equal(job.usage!.inputTokens, 10);
+				assert.equal(job.usage!.outputTokens, 5);
+			} finally {
+				researchModelOverride.current = undefined;
+			}
+		} finally {
+			restoreFetch();
+		}
+	});
+
+	it("aborting a delivering job drops the queued result steer and wakes with the abort note instead", async () => {
+		const chat = chatStore.createChat("Test: abort delivering job");
+		trackChat(chat.id);
+		chatStore.saveMessages(chat.id, [{ id: "u1", role: "user", parts: [{ type: "text", text: "hello" }] }] as never);
+		const job = jobStore.createJob({ kind: "research", chatId: chat.id, toolCallId: "call-abort3" });
+		trackJob(job.id);
+		const result: ResearchResult = { brief: "Done: it works.", sources: [{ title: "s", url: "https://x.test" }] };
+		jobStore.updateJob(job.id, { result, activity: "delivering result" });
+
+		// Block the pump's idle gates so the result steer sits QUEUED, not in flight.
+		assert.ok(beginTurn(chat.id));
+		const fired: SteerMessage[] = [];
+		setTurnRunner(async (_cid, messages) => {
+			fired.push(messages[messages.length - 1] as SteerMessage);
+			return 200;
+		});
+		void deliverJobResult(buildSubagentSteer(job, result)).catch(() => {
+			/* dropped by the abort below */
+		});
+		await waitFor(() => jobStore.getJob(job.id)!.activity === "delivering result", 8000, "queued");
+
+		const outcome = abortJob(job.id);
+		assert.ok(outcome.ok);
+		assert.equal(jobStore.getJob(job.id)!.status, "error");
+		assert.equal(jobStore.getJob(job.id)!.error, "Aborted by learner");
+
+		endTurn(chat.id); // gates open — whatever is left in the queue fires
+		await waitFor(() => fired.length > 0, 8000, "wake fired");
+		assert.equal(fired.length, 1, "exactly one delivery — the result steer was dropped");
+		assert.equal(fired[0]!.metadata.subagentResult.aborted, true, "the abort steer, not the result");
+		assert.equal(jobStore.getJob(job.id)!.status, "error", "terminal after the wake");
+	});
+
+	it("outcome reasons: not_found for unknown, not_running for settled jobs", () => {
+		assert.deepEqual(abortJob("job_does_not_exist"), { ok: false, reason: "not_found" });
+		const chat = chatStore.createChat("Test: abort settled job");
+		trackChat(chat.id);
+		const job = jobStore.createJob({ kind: "research", chatId: chat.id, toolCallId: "call-abort4" });
+		trackJob(job.id);
+		jobStore.updateJob(job.id, { status: "done", activity: "delivered" });
+		assert.deepEqual(abortJob(job.id), { ok: false, reason: "not_running" });
+	});
+});
@@ -12,7 +12,7 @@ import * as fs from "node:fs";
 import { chatStore } from "../src/store/chats.ts";
 import { PATHS } from "../config.ts";
 import { callCost, contextWindow, jobCost, sessionCost, sumSteps, turnCost } from "../src/core/prices.ts";
-import type { CallUsage, JobRecord, TurnUsage } from "../src/core/contracts.ts";
+import type { CallUsage, ExtractionRunUsage, JobRecord, TurnUsage } from "../src/core/contracts.ts";

 const call = (over: Partial<CallUsage> = {}): CallUsage => ({
 	inputTokens: 0,
@@ -134,3 +134,34 @@ test("appendTurnUsage: appends, is idempotent per messageId, survives saveMessag

 	fs.unlinkSync(`${PATHS.chats}/${record.id}.json`);
 });
+
+test("appendExtractionUsage: appends, idempotent per startedAt, not in list payload", () => {
+	const record = chatStore.createChat("extraction usage test chat");
+	const runEntry = (startedAt: number): ExtractionRunUsage => ({
+		startedAt,
+		at: "2026-09-09T00:00:00Z",
+		durationMs: 42_000,
+		stats: { nodes: 1, edges: 2, events: 3 },
+		stages: {
+			extractor: { model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 100, outputTokens: 10 })] },
+			auditor: { model: "z-ai/glm-5.3-flash", steps: [call({ inputTokens: 50, outputTokens: 5 })] },
+		},
+	});
+	chatStore.appendExtractionUsage(record.id, runEntry(1));
+	chatStore.appendExtractionUsage(record.id, runEntry(1)); // same run appended twice → no-op
+	chatStore.appendExtractionUsage(record.id, runEntry(2)); // a retried run is a NEW spend
+	let reread = chatStore.getChat(record.id);
+	assert.equal(reread?.extractions?.length, 2);
+	// failed runs (stats absent) persist with usage — money was spent
+	chatStore.appendExtractionUsage(record.id, { ...runEntry(3), stats: undefined });
+	reread = chatStore.getChat(record.id);
+	assert.equal(reread?.extractions?.length, 3);
+	assert.equal(reread?.extractions?.at(-1)?.stats, undefined);
+
+	// the list payload must NOT leak the extractions log
+	const listed = chatStore.listChats().find((c) => c.id === record.id);
+	assert.ok(listed);
+	assert.equal("extractions" in listed, false);
+
+	fs.unlinkSync(`${PATHS.chats}/${record.id}.json`);
+});