[2] Keep work tied to its account and task

src/api/organizations.test.ts

@@ -16,7 +16,10 @@ function jsonResponse(body: unknown, status = 200): Response {
 describe("verifyApiKey", () => {
 	it("returns the org id the key is scoped to", async () => {
 		const fetchImpl = vi.fn(async () => jsonResponse({ organizationId: "org-1", userId: "user-1" }))
-		await expect(verifyApiKey("key", { fetch: fetchImpl })).resolves.toEqual({ organizationId: "org-1" })
+		await expect(verifyApiKey("key", { fetch: fetchImpl })).resolves.toEqual({
+			organizationId: "org-1",
+			userId: "user-1",
+		})
 		expect(fetchImpl).toHaveBeenCalledWith("https://api.test/ai-optimizer/v1beta/api-keys:verify", {
 			method: "POST",
 			headers: { Authorization: "Bearer key", Accept: "application/json" },

src/api/organizations.ts

@@ -3,6 +3,7 @@ import { fetchWithRetry } from "../utils/http.js"
 
 export interface VerifyApiKeyResponse {
 	organizationId: string
+	userId?: string
 }
 
 export interface Organization {
@@ -22,9 +23,14 @@ function resolveEndpoint(): string {
  */
 export async function verifyApiKey(
 	apiKey: string,
-	options?: { fetch?: typeof globalThis.fetch },
+	options?: {
+		fetch?: typeof globalThis.fetch
+		endpoint?: string
+		signal?: AbortSignal
+		retry?: { maxRetries: number }
+	},
 ): Promise<VerifyApiKeyResponse> {
-	const endpoint = resolveEndpoint()
+	const endpoint = options?.endpoint ?? resolveEndpoint()
 	const fetchImpl = options?.fetch ?? globalThis.fetch
 
 	const url = `${endpoint}/ai-optimizer/v1beta/api-keys:verify`
@@ -37,7 +43,7 @@ export async function verifyApiKey(
 				Accept: "application/json",
 			},
 		},
-		{ fetchImpl },
+		{ fetchImpl, signal: options?.signal, retry: options?.retry },
 	)
 
 	if (!resp.ok) {
@@ -52,7 +58,10 @@ export async function verifyApiKey(
 		throw new Error(`Missing organizationId in api-keys:verify response from ${url}`)
 	}
 
-	return { organizationId: data.organizationId }
+	return {
+		organizationId: data.organizationId,
+		...(typeof data.userId === "string" && data.userId ? { userId: data.userId } : {}),
+	}
 }
 
 /** Fetch organization by id. */

src/cli.ts

@@ -159,6 +159,7 @@ import traceIdExtension from "./extensions/trace-id.js"
 import uiExtension from "./extensions/ui.js"
 import webFetchExtension from "./extensions/web-fetch/index.js"
 import webSearchExtension from "./extensions/web-search/index.js"
+import { COST_PER_PR_RESOURCE_ID } from "./extensions/work-attribution/resource.js"
 import { createWorkAttributionExtension } from "./extensions/work-attribution.js"
 import { normalizeAtFileArgs } from "./fs-paths.js"
 import { installGlobalFetchInstrumentation } from "./http/instrument-fetch.js"
@@ -827,8 +828,10 @@ try {
 			traceIdExtension,
 			contextAssemblyExtension,
 			cacheSummaryExtension,
-			createWorkAttributionExtension(),
-			pullRequestStatusExtension,
+			...enabledExtensionFactories([
+				{ id: COST_PER_PR_RESOURCE_ID, factory: createWorkAttributionExtension() },
+				{ id: COST_PER_PR_RESOURCE_ID, factory: pullRequestStatusExtension },
+			] satisfies ManagedExtensionFactory[]),
 			requestTimingExtension,
 			llmResponseLogExtension,
 			activityExtension,

src/config/json.ts

@@ -1,5 +1,7 @@
+import { randomUUID } from "node:crypto"
 import { mkdirSync, readFileSync, renameSync, statSync, writeFileSync } from "node:fs"
-import { dirname, resolve } from "node:path"
+import { open, rename, rm } from "node:fs/promises"
+import { basename, dirname, join, resolve } from "node:path"
 
 /**
  * The `.jsonc` sibling tried when a `.json` path is absent. Single source of
@@ -175,6 +177,27 @@ export async function writeJsonAsync(path: string, data: unknown): Promise<void>
 	})
 }
 
+/**
+ * Crash-safe replacement: syncs a private temporary file, lets the caller confirm
+ * it still owns the write, then renames it over `path`.
+ */
+export async function writeFileDurably(path: string, data: string, beforeRename = () => {}): Promise<void> {
+	const temporary = join(dirname(path), `.${basename(path)}-${randomUUID()}.tmp`)
+	try {
+		const file = await open(temporary, "wx", 0o600)
+		try {
+			await file.writeFile(data)
+			await file.sync()
+		} finally {
+			await file.close()
+		}
+		beforeRename()
+		await rename(temporary, path)
+	} finally {
+		await rm(temporary, { force: true })
+	}
+}
+
 /** Atomic raw write, used by the OpenClaw .env writer. */
 export function writeFileAtomic(path: string, data: string | Uint8Array): void {
 	mkdirSync(dirname(path), { recursive: true })

src/extensions/__mocks__/context.ts

@@ -9,6 +9,7 @@ import type {
 	TerminalInputHandler,
 } from "@earendil-works/pi-coding-agent"
 import { type Mocked, vi } from "vitest"
+import { createModelRegistry } from "./model-registry.js"
 
 export function createCommandContext(): ExtensionCommandContext {
 	return {
@@ -31,6 +32,46 @@ export function sendTerminalInput(ctx: ExtensionContext, data: string): void {
 	}
 }
 
+/**
+ * {@link createContext} built like Pi's runner context: each property is a getter read on access, and every read
+ * throws once `invalidate()` marks the context stale. Use it for code that keeps or copies `ctx`, where plain data
+ * properties would hide a frozen copy. Pi checks its methods when called; this mock checks them when read.
+ */
+export function createLiveContext(overrides?: Parameters<typeof createContext>[0]): {
+	ctx: ExtensionContext
+	/** Changes what later reads return, as a model switch does in Pi. */
+	set<K extends keyof ExtensionContext>(key: K, value: ExtensionContext[K]): void
+	/** Marks the context stale, as Pi does after a session replacement or reload. */
+	invalidate(): void
+} {
+	const values: Record<string, unknown> = { ...createContext(overrides) }
+	let stale = false
+	const ctx = Object.defineProperties(
+		{},
+		Object.fromEntries(
+			Object.keys(values).map((key) => [
+				key,
+				{
+					enumerable: true,
+					get: () => {
+						if (stale) throw new Error("This extension ctx is stale after session replacement or reload.")
+						return values[key]
+					},
+				},
+			]),
+		),
+	) as ExtensionContext
+	return {
+		ctx,
+		set: (key, value) => {
+			values[key] = value
+		},
+		invalidate: () => {
+			stale = true
+		},
+	}
+}
+
 export function createContext(
 	overrides?: Mocked<
 		Omit<Partial<ExtensionContext>, "ui" | "sessionManager" | "modelRegistry" | "model" | "getContextUsage"> & {
@@ -53,7 +94,9 @@ export function createContext(
 		cwd: "/tmp",
 		scopedModels: [],
 		isIdle: vi.fn(),
+		hasPendingMessages: vi.fn(() => false),
 		getContextUsage: vi.fn().mockReturnValue(undefined),
+		modelRegistry: createModelRegistry(),
 		...overrides,
 		model: overrides?.model
 			? {

src/extensions/__mocks__/work-scope.ts

@@ -0,0 +1,15 @@
+import type { WorkScopeSnapshot } from "../work-attribution/scope.js"
+
+export function createWorkScopeSnapshot(repository = "/project/.git"): WorkScopeSnapshot {
+	return {
+		scope: {
+			repository,
+			account: {
+				apiUrl: "https://account.example/api",
+				organizationId: "30000000-0000-4000-8000-000000000003",
+				userId: "40000000-0000-4000-8000-000000000004",
+			},
+		},
+		isCurrent: () => true,
+	}
+}

src/extensions/agents/manager/agent-manager.queued-work.test.ts

@@ -0,0 +1,128 @@
+import { randomUUID } from "node:crypto"
+import { mkdtempSync, rmSync } from "node:fs"
+import { tmpdir } from "node:os"
+import { join } from "node:path"
+import { afterEach, beforeEach, expect, it, vi } from "vitest"
+
+vi.mock("./agent-runner.js", () => ({
+	runAgent: vi.fn(),
+	resumeAgent: vi.fn(),
+	MIN_TOKEN_BUDGET: 1024,
+	MIN_FINALIZE_TOKEN_BUDGET: 256,
+}))
+vi.mock("./remote-agent-runner.js", () => ({
+	runRemoteAgent: vi.fn(),
+	continueRemoteAgent: vi.fn(),
+	attachRemoteAgent: vi.fn(),
+	isRemoteSessionConnected: vi.fn(),
+}))
+
+import type { AgentSession } from "@earendil-works/pi-coding-agent"
+import { createContext, createLiveContext } from "../../__mocks__/context.js"
+import { createExtensionApi } from "../../__mocks__/extension-api.js"
+import { createModel } from "../../__mocks__/model-registry.js"
+import { flushWorkSummaries } from "../../work-attribution/summary.js"
+import { appendWorkRecord, getWorkId, getWorkSegment, setWorkId } from "../../work-attribution.js"
+import { AgentManager } from "./agent-manager.js"
+import { runAgent } from "./agent-runner.js"
+
+let dir: string
+beforeEach(() => {
+	dir = mkdtempSync(join(tmpdir(), "kimchi-agent-queued-"))
+	vi.stubEnv("PI_CODING_AGENT_DIR", dir)
+})
+afterEach(async () => {
+	await flushWorkSummaries()
+	vi.unstubAllEnvs()
+	vi.clearAllMocks()
+	rmSync(dir, { recursive: true, force: true })
+})
+
+it("starts a queued background child in the work and segment active when it was spawned", async () => {
+	const ctx = createContext({ cwd: dir })
+	const spawnedIn = getWorkId(ctx)
+	const spawnedSegment = getWorkSegment(ctx)
+	// runAgentInner() captures getWorkId(ctx)/getWorkSegment(ctx) as soon as runAgent() starts.
+	const inherited: { workId: string; segment: unknown }[] = []
+	let finishFirst!: () => void
+	vi.mocked(runAgent).mockImplementation(async (childCtx) => {
+		inherited.push({ workId: getWorkId(childCtx), segment: getWorkSegment(childCtx) })
+		if (inherited.length === 1) await new Promise<void>((resolve) => (finishFirst = resolve))
+		return {
+			responseText: "done",
+			session: { dispose: vi.fn() } as unknown as AgentSession,
+			aborted: false,
+			steered: false,
+		}
+	})
+	const manager = new AgentManager(undefined, 1)
+	const pi = createExtensionApi().api
+	try {
+		manager.spawn(pi, ctx, "Explore", "one", { description: "one", isBackground: true })
+		const queued = manager.spawn(pi, ctx, "Explore", "two", { description: "two", isBackground: true })
+		expect(manager.getRecord(queued)?.status).toBe("queued")
+		// The parent moves on (/work new, a pasted plan, or a semantic "new" decision).
+		setWorkId(ctx)
+		finishFirst()
+		await vi.waitFor(() => expect(inherited).toHaveLength(2))
+		expect(inherited[1]).toEqual({ workId: spawnedIn, segment: spawnedSegment })
+	} finally {
+		manager.dispose()
+	}
+})
+
+it("pins only the work for a queued child; Pi's context getters and stale guard stay live", async () => {
+	const parent = createLiveContext({ cwd: dir, model: createModel("model-at-spawn") })
+	const spawnedIn = randomUUID()
+	const spawnedSegment = { id: randomUUID(), attribution: "session", reason: "matching-disabled" } as const
+	// A restored ledger gives the parent an active input segment.
+	appendWorkRecord(parent.ctx, { type: "work", segment: spawnedSegment }, spawnedIn)
+	expect([getWorkId(parent.ctx), getWorkSegment(parent.ctx)]).toEqual([spawnedIn, spawnedSegment])
+	const read = (value: () => unknown) => {
+		try {
+			return value()
+		} catch (error) {
+			return error instanceof Error ? error.message : String(error)
+		}
+	}
+	const started: Record<string, unknown>[] = []
+	const finish: (() => void)[] = []
+	vi.mocked(runAgent).mockImplementation(async (childCtx) => {
+		started.push({
+			workId: getWorkId(childCtx),
+			segment: getWorkSegment(childCtx),
+			model: read(() => childCtx.model?.id),
+			session: read(() => childCtx.sessionManager.getSessionId()),
+		})
+		await new Promise<void>((resolve) => finish.push(resolve))
+		return {
+			responseText: "done",
+			session: { dispose: vi.fn() } as unknown as AgentSession,
+			aborted: false,
+			steered: false,
+		}
+	})
+	const manager = new AgentManager(undefined, 1)
+	const pi = createExtensionApi().api
+	try {
+		for (const name of ["running", "after-model-change", "after-session-replacement"])
+			manager.spawn(pi, parent.ctx, "Explore", name, { description: name, isBackground: true })
+		setWorkId(parent.ctx)
+		parent.set("model", createModel("model-at-start"))
+		finish[0]()
+		await vi.waitFor(() => expect(started).toHaveLength(2))
+		parent.invalidate()
+		finish[1]()
+		await vi.waitFor(() => expect(started).toHaveLength(3))
+		finish[2]()
+		const pinned = { workId: spawnedIn, segment: spawnedSegment }
+		expect(started[1]).toEqual({ ...pinned, model: "model-at-start", session: "test-session" })
+		expect(started[2]).toEqual({
+			...pinned,
+			model: expect.stringContaining("stale"),
+			session: expect.stringContaining("stale"),
+		})
+	} finally {
+		manager.dispose()
+	}
+})

src/extensions/agents/manager/agent-manager.test.ts

@@ -59,9 +59,12 @@ import type { AgentSession, ExtensionAPI, ExtensionContext } from "@earendil-wor
 import { loadWorkspaceFile, WorkspaceFileError } from "../../../sandbox/cloud/workspace-file.js"
 import { listWorkspaces } from "../../../sandbox/cloud/workspaces.js"
 import { SESSION_TAG_PARENT_SESSION_ID } from "../../../sandbox/worker/types.js"
+import { createContext } from "../../__mocks__/context.js"
+import { createExtensionApi } from "../../__mocks__/extension-api.js"
 import { captureBaseline, resolveSandboxGitConnection, SandboxGitError } from "../../remote-run/sandbox-git.js"
 import { resolveClonePlan } from "../../teleport/provisioning/clone-plan.js"
 import { resolveGitToken } from "../../teleport/provisioning/git-token.js"
+import { flushWorkSummaries } from "../../work-attribution/summary.js"
 import type { AgentRecord } from "../personas/types.js"
 import type { PersistedGitWorkflow, RemoteRunState } from "../remote-run-persistence.js"
 import { AgentManager, buildAgentOutcome } from "./agent-manager.js"
@@ -87,12 +90,16 @@ const mockLoadWorkspaceFile = vi.mocked(loadWorkspaceFile)
 const mockResolveSandboxGitConnection = vi.mocked(resolveSandboxGitConnection)
 const mockCaptureBaseline = vi.mocked(captureBaseline)
 
+afterEach(async () => {
+	await flushWorkSummaries()
+})
+
 function fakePi(): ExtensionAPI {
-	return {} as ExtensionAPI
+	return createExtensionApi().api
 }
 
 function fakeCtx(): ExtensionContext {
-	return {} as ExtensionContext
+	return createContext()
 }
 
 describe("AgentManager", () => {
@@ -930,11 +937,11 @@ describe("AgentManager visibility", () => {
 	it("stores system visibility on queued records", () => {
 		const manager = new AgentManager(undefined, 0)
 		try {
-			const first = manager.spawn({} as never, {} as never, "General-Purpose", "one", {
+			const first = manager.spawn(fakePi(), fakeCtx(), "General-Purpose", "one", {
 				description: "visible agent",
 				isBackground: true,
 			})
-			const second = manager.spawn({} as never, {} as never, "General-Purpose", "two", {
+			const second = manager.spawn(fakePi(), fakeCtx(), "General-Purpose", "two", {
 				description: "system agent",
 				isBackground: true,
 				visibility: "system",