diff --git a/src/features/tool-metadata-store/index.ts b/src/features/tool-metadata-store/index.ts index f9c4e28ad..cc89a9f72 100644 --- a/src/features/tool-metadata-store/index.ts +++ b/src/features/tool-metadata-store/index.ts @@ -5,3 +5,10 @@ export { storeToolMetadata, } from "./store" export type { PendingToolMetadata } from "./store" +export { resolveToolCallID } from "./resolve-tool-call-id" +export type { ToolCallIDCarrier } from "./resolve-tool-call-id" +export { buildTaskMetadataBlock, extractTaskLink, parseTaskMetadataBlock } from "./task-metadata-contract" +export type { TaskLink } from "./task-metadata-contract" +export { publishToolMetadata } from "./publish-tool-metadata" +export { recoverToolMetadata } from "./recover-tool-metadata" +export type { ToolMetadataPublisherContext } from "./publish-tool-metadata" diff --git a/src/features/tool-metadata-store/integration.test.ts b/src/features/tool-metadata-store/integration.test.ts new file mode 100644 index 000000000..2604a5edf --- /dev/null +++ b/src/features/tool-metadata-store/integration.test.ts @@ -0,0 +1,96 @@ +import { beforeEach, describe, expect, test } from "bun:test" + +import { clearPendingStore, getPendingStoreSize } from "./store" +import { publishToolMetadata } from "./publish-tool-metadata" +import { recoverToolMetadata } from "./recover-tool-metadata" + +describe("tool-metadata-store integration", () => { + beforeEach(() => { + clearPendingStore() + }) + + test("#given stored metadata #when publishing then recovering #then the round trip preserves the payload", async () => { + // given + const payload = { title: "Task", metadata: { sessionId: "ses_child" } } + + // when + await publishToolMetadata({ sessionID: "ses_parent", callID: "call_123" }, payload) + const recovered = recoverToolMetadata("ses_parent", { callID: "call_123" }) + + // then + expect(recovered).toEqual(payload) + }) + + test("#given call id casing mismatch #when publishing and recovering #then canonical resolution still matches", async () => { + // given + const payload = { title: "Task", metadata: { sessionId: "ses_child" } } + + // when + await publishToolMetadata({ sessionID: "ses_parent", callId: "call_case" }, payload) + const recovered = recoverToolMetadata("ses_parent", { callID: "call_case" }) + + // then + expect(recovered).toEqual(payload) + }) + + test("#given blank call id #when publishing #then nothing is stored", async () => { + // given + const payload = { title: "Task" } + + // when + const result = await publishToolMetadata({ sessionID: "ses_parent", callID: " " }, payload) + const recovered = recoverToolMetadata("ses_parent", { callID: "call_blank" }) + + // then + expect(result).toEqual({ stored: false }) + expect(recovered).toBeUndefined() + expect(getPendingStoreSize()).toBe(0) + }) + + test("#given missing call id #when publishing #then nothing is stored", async () => { + // given + const payload = { title: "Task" } + + // when + const result = await publishToolMetadata({ sessionID: "ses_parent" }, payload) + + // then + expect(result).toEqual({ stored: false }) + expect(getPendingStoreSize()).toBe(0) + }) + + test("#given same session with different call ids #when publishing twice #then each entry stays isolated", async () => { + // given + await publishToolMetadata({ sessionID: "ses_parent", callID: "call_a" }, { title: "A" }) + await publishToolMetadata({ sessionID: "ses_parent", callID: "call_b" }, { title: "B" }) + + // when + const recoveredA = recoverToolMetadata("ses_parent", { callID: "call_a" }) + const recoveredB = recoverToolMetadata("ses_parent", { callID: "call_b" }) + + // then + expect(recoveredA).toEqual({ title: "A" }) + expect(recoveredB).toEqual({ title: "B" }) + }) + + test("#given stale metadata #when a fresh entry is stored after the timeout #then stale entries are cleaned up", async () => { + // given + const originalDateNow = Date.now + let now = 0 + Date.now = () => now + + try { + await publishToolMetadata({ sessionID: "ses_parent", callID: "call_old" }, { title: "Old" }) + now = 15 * 60 * 1000 + 1 + + // when + await publishToolMetadata({ sessionID: "ses_parent", callID: "call_new" }, { title: "New" }) + + // then + expect(recoverToolMetadata("ses_parent", { callID: "call_old" })).toBeUndefined() + expect(recoverToolMetadata("ses_parent", { callID: "call_new" })).toEqual({ title: "New" }) + } finally { + Date.now = originalDateNow + } + }) +}) diff --git a/src/features/tool-metadata-store/publish-tool-metadata.test.ts b/src/features/tool-metadata-store/publish-tool-metadata.test.ts new file mode 100644 index 000000000..704fe2392 --- /dev/null +++ b/src/features/tool-metadata-store/publish-tool-metadata.test.ts @@ -0,0 +1,76 @@ +import { beforeEach, describe, expect, test } from "bun:test" + +import { clearPendingStore, consumeToolMetadata } from "./store" +import { publishToolMetadata } from "./publish-tool-metadata" + +describe("publishToolMetadata", () => { + beforeEach(() => { + clearPendingStore() + }) + + test("#given metadata context and call id #when publishing #then it awaits metadata and stores the payload", async () => { + // given + const calls: string[] = [] + let metadataFinished = false + const payload = { title: "Task", metadata: { sessionId: "ses_child" } } + + // when + const result = await publishToolMetadata( + { + sessionID: "ses_parent", + callID: "call_123", + metadata: async input => { + calls.push(input.title ?? "") + await new Promise(resolve => setTimeout(resolve, 1)) + metadataFinished = true + }, + }, + payload + ) + + // then + expect(result).toEqual({ stored: true }) + expect(metadataFinished).toBe(true) + expect(calls).toEqual(["Task"]) + expect(consumeToolMetadata("ses_parent", "call_123")).toEqual(payload) + }) + + test("#given legacy call id variant #when publishing #then it stores with the canonical resolver", async () => { + // given + const payload = { title: "Task", metadata: { sessionId: "ses_child" } } + + // when + const result = await publishToolMetadata( + { + sessionID: "ses_parent", + callId: " call_legacy ", + }, + payload + ) + + // then + expect(result).toEqual({ stored: true }) + expect(consumeToolMetadata("ses_parent", "call_legacy")).toEqual(payload) + }) + + test("#given missing call id #when publishing #then it still emits metadata but skips storing", async () => { + // given + let metadataCalls = 0 + + // when + const result = await publishToolMetadata( + { + sessionID: "ses_parent", + metadata: () => { + metadataCalls += 1 + }, + }, + { title: "Task" } + ) + + // then + expect(result).toEqual({ stored: false }) + expect(metadataCalls).toBe(1) + expect(consumeToolMetadata("ses_parent", "call_missing")).toBeUndefined() + }) +}) diff --git a/src/features/tool-metadata-store/publish-tool-metadata.ts b/src/features/tool-metadata-store/publish-tool-metadata.ts new file mode 100644 index 000000000..84e69155e --- /dev/null +++ b/src/features/tool-metadata-store/publish-tool-metadata.ts @@ -0,0 +1,28 @@ +import { log } from "../../shared/logger" +import { resolveToolCallID, type ToolCallIDCarrier } from "./resolve-tool-call-id" +import { storeToolMetadata, type PendingToolMetadata } from "./store" + +export interface ToolMetadataPublisherContext extends ToolCallIDCarrier { + sessionID: string + metadata?: (input: PendingToolMetadata) => void | Promise +} + +export async function publishToolMetadata( + ctx: ToolMetadataPublisherContext, + payload: PendingToolMetadata +): Promise<{ stored: boolean }> { + await ctx.metadata?.(payload) + + const callID = resolveToolCallID(ctx) + if (!callID) { + log("[tool-metadata-store] Skipping metadata store publish because tool call ID is unavailable", { + sessionID: ctx.sessionID, + hasTitle: typeof payload.title === "string", + hasMetadata: payload.metadata !== undefined, + }) + return { stored: false } + } + + storeToolMetadata(ctx.sessionID, callID, payload) + return { stored: true } +} diff --git a/src/features/tool-metadata-store/recover-tool-metadata.test.ts b/src/features/tool-metadata-store/recover-tool-metadata.test.ts new file mode 100644 index 000000000..e97e87eea --- /dev/null +++ b/src/features/tool-metadata-store/recover-tool-metadata.test.ts @@ -0,0 +1,47 @@ +import { beforeEach, describe, expect, test } from "bun:test" + +import { recoverToolMetadata } from "./recover-tool-metadata" +import { clearPendingStore, storeToolMetadata } from "./store" + +describe("recoverToolMetadata", () => { + beforeEach(() => { + clearPendingStore() + }) + + test("#given stored metadata and call id variant #when recovering #then it finds the stored payload", () => { + // given + const payload = { title: "Recovered", metadata: { sessionId: "ses_child" } } + storeToolMetadata("ses_parent", "call_123", payload) + + // when + const recovered = recoverToolMetadata("ses_parent", { callId: " call_123 " }) + + // then + expect(recovered).toEqual(payload) + }) + + test("#given direct string call id #when recovering #then it consumes the stored payload", () => { + // given + const payload = { title: "Recovered" } + storeToolMetadata("ses_parent", "call_456", payload) + + // when + const recovered = recoverToolMetadata("ses_parent", "call_456") + + // then + expect(recovered).toEqual(payload) + }) + + test("#given missing or blank call id #when recovering #then it returns undefined", () => { + // given + storeToolMetadata("ses_parent", "call_789", { title: "Recovered" }) + + // when + const missing = recoverToolMetadata("ses_parent", undefined) + const blank = recoverToolMetadata("ses_parent", { callID: " " }) + + // then + expect(missing).toBeUndefined() + expect(blank).toBeUndefined() + }) +}) diff --git a/src/features/tool-metadata-store/recover-tool-metadata.ts b/src/features/tool-metadata-store/recover-tool-metadata.ts new file mode 100644 index 000000000..17a2278e6 --- /dev/null +++ b/src/features/tool-metadata-store/recover-tool-metadata.ts @@ -0,0 +1,18 @@ +import { consumeToolMetadata, type PendingToolMetadata } from "./store" +import { resolveToolCallID, type ToolCallIDCarrier } from "./resolve-tool-call-id" + +export function recoverToolMetadata( + sessionID: string, + source: ToolCallIDCarrier | string | undefined +): PendingToolMetadata | undefined { + if (typeof source === "string") { + return consumeToolMetadata(sessionID, source) + } + + const callID = source ? resolveToolCallID(source) : undefined + if (!callID) { + return undefined + } + + return consumeToolMetadata(sessionID, callID) +} diff --git a/src/features/tool-metadata-store/resolve-tool-call-id.test.ts b/src/features/tool-metadata-store/resolve-tool-call-id.test.ts new file mode 100644 index 000000000..2ad63a516 --- /dev/null +++ b/src/features/tool-metadata-store/resolve-tool-call-id.test.ts @@ -0,0 +1,88 @@ +import { describe, expect, test } from "bun:test" + +import { resolveToolCallID, type ToolCallIDCarrier } from "./resolve-tool-call-id" + +describe("resolveToolCallID", () => { + function makeCtx(overrides: Partial = {}): ToolCallIDCarrier { + return { + ...overrides, + } + } + + test("#given callID is set #when resolving #then it returns callID", () => { + // given + const ctx = makeCtx({ callID: "call_abc" }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBe("call_abc") + }) + + test("#given only callId is set #when resolving #then it returns callId", () => { + // given + const ctx = makeCtx({ callId: "call_def" }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBe("call_def") + }) + + test("#given only call_id is set #when resolving #then it returns call_id", () => { + // given + const ctx = makeCtx({ call_id: "call_ghi" }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBe("call_ghi") + }) + + test("#given surrounding whitespace #when resolving #then it trims the value", () => { + // given + const ctx = makeCtx({ callID: " call_trimmed " }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBe("call_trimmed") + }) + + test("#given blank callID #when resolving #then it returns undefined", () => { + // given + const ctx = makeCtx({ callID: "" }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBeUndefined() + }) + + test("#given whitespace callID #when resolving #then it returns undefined", () => { + // given + const ctx = makeCtx({ callID: " " }) + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBeUndefined() + }) + + test("#given no call id variants #when resolving #then it returns undefined", () => { + // given + const ctx = makeCtx() + + // when + const result = resolveToolCallID(ctx) + + // then + expect(result).toBeUndefined() + }) +}) diff --git a/src/features/tool-metadata-store/resolve-tool-call-id.ts b/src/features/tool-metadata-store/resolve-tool-call-id.ts new file mode 100644 index 000000000..d5879ae5f --- /dev/null +++ b/src/features/tool-metadata-store/resolve-tool-call-id.ts @@ -0,0 +1,26 @@ +import { log } from "../../shared/logger" + +export interface ToolCallIDCarrier { + callID?: string + callId?: string + call_id?: string +} + +function normalizeCallID(value: unknown): string | undefined { + if (typeof value !== "string") { + return undefined + } + + const trimmed = value.trim() + return trimmed === "" ? undefined : trimmed +} + +export function resolveToolCallID(ctx: ToolCallIDCarrier): string | undefined { + const resolved = normalizeCallID(ctx.callID) ?? normalizeCallID(ctx.callId) ?? normalizeCallID(ctx.call_id) + + if (!resolved) { + log("[tool-metadata-store] Missing tool call ID for metadata correlation") + } + + return resolved +} diff --git a/src/features/tool-metadata-store/task-metadata-contract.test.ts b/src/features/tool-metadata-store/task-metadata-contract.test.ts new file mode 100644 index 000000000..2cf2aaa7c --- /dev/null +++ b/src/features/tool-metadata-store/task-metadata-contract.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, test } from "bun:test" + +import { buildTaskMetadataBlock, extractTaskLink, parseTaskMetadataBlock } from "./task-metadata-contract" + +describe("buildTaskMetadataBlock", () => { + test("#given only session id #when building #then it preserves the frozen block format", () => { + // given + const link = { sessionId: "ses_abc" } + + // when + const block = buildTaskMetadataBlock(link) + + // then + expect(block).toBe("\nsession_id: ses_abc\n") + }) + + test("#given extended task metadata #when building #then it emits optional lines in order", () => { + // given + const link = { + sessionId: "ses_bg_123", + taskId: "bg_123", + backgroundTaskId: "bg_123", + agent: "explore", + category: "quick", + } + + // when + const block = buildTaskMetadataBlock(link) + + // then + expect(block).toBe( + "\nsession_id: ses_bg_123\ntask_id: bg_123\nbackground_task_id: bg_123\nsubagent: explore\ncategory: quick\n" + ) + }) +}) + +describe("parseTaskMetadataBlock", () => { + test("#given a task metadata block #when parsing #then it extracts the structured link", () => { + // given + const text = "\nsession_id: ses_sync_123\ntask_id: task_123\nbackground_task_id: bg_123\nsubagent: oracle\ncategory: deep\n" + + // when + const parsed = parseTaskMetadataBlock(text) + + // then + expect(parsed).toEqual({ + sessionId: "ses_sync_123", + taskId: "task_123", + backgroundTaskId: "bg_123", + agent: "oracle", + category: "deep", + }) + }) + + test("#given text without metadata #when parsing #then it returns an empty link", () => { + // given + const text = "Task completed without metadata" + + // when + const parsed = parseTaskMetadataBlock(text) + + // then + expect(parsed).toEqual({}) + }) +}) + +describe("extractTaskLink", () => { + test("#given metadata session aliases #when extracting #then metadata wins over output text", () => { + // given + const metadata = { + sessionID: "ses_meta_123", + task_id: "task_meta_123", + background_task_id: "bg_meta_123", + subagent: "atlas", + category: "unspecified-high", + } + const output = "\nsession_id: ses_text_456\n" + + // when + const extracted = extractTaskLink(metadata, output) + + // then + expect(extracted).toEqual({ + sessionId: "ses_meta_123", + taskId: "task_meta_123", + backgroundTaskId: "bg_meta_123", + agent: "atlas", + category: "unspecified-high", + }) + }) + + test("#given missing metadata #when extracting #then it falls back to task metadata text", () => { + // given + const output = "Task completed.\n\n\nsession_id: ses_text_456\nsubagent: oracle\n" + + // when + const extracted = extractTaskLink(undefined, output) + + // then + expect(extracted).toEqual({ + sessionId: "ses_text_456", + agent: "oracle", + }) + }) + + test("#given explicit session id output #when extracting #then it preserves Session ID compatibility", () => { + // given + const output = "Background task launched.\n\nSession ID: ses_bg_789" + + // when + const extracted = extractTaskLink(undefined, output) + + // then + expect(extracted).toEqual({ sessionId: "ses_bg_789" }) + }) +}) diff --git a/src/features/tool-metadata-store/task-metadata-contract.ts b/src/features/tool-metadata-store/task-metadata-contract.ts new file mode 100644 index 000000000..4044482e5 --- /dev/null +++ b/src/features/tool-metadata-store/task-metadata-contract.ts @@ -0,0 +1,144 @@ +import { log } from "../../shared/logger" + +export interface TaskLink { + sessionId?: string + taskId?: string + backgroundTaskId?: string + agent?: string + category?: string +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null +} + +function readString(value: unknown): string | undefined { + if (typeof value !== "string") { + return undefined + } + + const trimmed = value.trim() + return trimmed === "" ? undefined : trimmed +} + +function readSessionIdFromMetadata(metadata: Record): string | undefined { + return readString(metadata.sessionId) ?? readString(metadata.sessionID) ?? readString(metadata.session_id) +} + +function readTaskIdFromMetadata(metadata: Record): string | undefined { + return readString(metadata.taskId) ?? readString(metadata.taskID) ?? readString(metadata.task_id) +} + +function readBackgroundTaskIdFromMetadata(metadata: Record): string | undefined { + return readString(metadata.backgroundTaskId) + ?? readString(metadata.backgroundTaskID) + ?? readString(metadata.background_task_id) +} + +function readAgentFromMetadata(metadata: Record): string | undefined { + return readString(metadata.agent) ?? readString(metadata.subagent) +} + +function readCategoryFromMetadata(metadata: Record): string | undefined { + return readString(metadata.category) +} + +function extractTaskMetadataContent(text: string): string | undefined { + const blocks = [...text.matchAll(/([\s\S]*?)<\/task_metadata>/gi)] + return blocks.at(-1)?.[1] +} + +function extractExplicitSessionId(text: string): string | undefined { + const matches = [...text.matchAll(/Session ID:\s*(ses_[a-zA-Z0-9_-]+)/g)] + return matches.at(-1)?.[1] +} + +export function buildTaskMetadataBlock(link: TaskLink): string { + const lines: string[] = [] + + if (link.sessionId) { + lines.push(`session_id: ${link.sessionId}`) + } + if (link.taskId) { + lines.push(`task_id: ${link.taskId}`) + } + if (link.backgroundTaskId) { + lines.push(`background_task_id: ${link.backgroundTaskId}`) + } + if (link.agent) { + lines.push(`subagent: ${link.agent}`) + } + if (link.category) { + lines.push(`category: ${link.category}`) + } + + return `\n${lines.join("\n")}\n` +} + +export function parseTaskMetadataBlock(text: string): TaskLink { + const blockContent = extractTaskMetadataContent(text) ?? text + const lines = blockContent + .split("\n") + .map(line => line.trim()) + .filter(Boolean) + + const parsed: TaskLink = {} + + for (const line of lines) { + const separatorIndex = line.indexOf(":") + if (separatorIndex === -1) { + continue + } + + const key = line.slice(0, separatorIndex).trim().toLowerCase() + const value = readString(line.slice(separatorIndex + 1)) + + if (!value) { + continue + } + + if (key === "session_id") { + parsed.sessionId = value + } else if (key === "task_id") { + parsed.taskId = value + } else if (key === "background_task_id") { + parsed.backgroundTaskId = value + } else if (key === "subagent" || key === "agent") { + parsed.agent = value + } else if (key === "category") { + parsed.category = value + } + } + + return parsed +} + +export function extractTaskLink(metadata: unknown, outputText: string): TaskLink { + if (isRecord(metadata)) { + const metadataLink: TaskLink = { + sessionId: readSessionIdFromMetadata(metadata), + taskId: readTaskIdFromMetadata(metadata), + backgroundTaskId: readBackgroundTaskIdFromMetadata(metadata), + agent: readAgentFromMetadata(metadata), + category: readCategoryFromMetadata(metadata), + } + + if (metadataLink.sessionId || metadataLink.taskId || metadataLink.backgroundTaskId || metadataLink.agent || metadataLink.category) { + return metadataLink + } + } + + const parsed = parseTaskMetadataBlock(outputText) + if (parsed.sessionId || parsed.taskId || parsed.backgroundTaskId || parsed.agent || parsed.category) { + log("[tool-metadata-store] Falling back to parsing") + return parsed + } + + const explicitSessionId = extractExplicitSessionId(outputText) + if (explicitSessionId) { + log("[tool-metadata-store] Falling back to explicit Session ID parsing") + return { sessionId: explicitSessionId } + } + + return {} +}