gsd-pi 2.80.0-dev.fbe7c8c6f → 2.81.0-dev.3cddbbba2

package/packages/pi-agent-core/src/token-audit.test.ts

@@ -0,0 +1,189 @@
+// Project/App: GSD-2
+// File Purpose: Tests for provider-boundary token payload audit helpers.
+
+import assert from "node:assert/strict";
+import test from "node:test";
+import { Type } from "@sinclair/typebox";
+import type { Context } from "@gsd/pi-ai";
+import type { AgentMessage } from "./types.js";
+import {
+	buildProviderPayloadAuditSummary,
+	buildTokenAuditSummary,
+	maybeLogProviderPayloadAudit,
+	maybeLogTokenAudit,
+} from "./token-audit.js";
+
+test("buildTokenAuditSummary reports payload section sizes without content fields", () => {
+	const context: Context = {
+		systemPrompt: "system prompt",
+		tools: [
+			{
+				name: "read",
+				description: "Read a file",
+				parameters: Type.Object({ path: Type.String() }),
+			},
+			{
+				name: "large_tool",
+				description: "Large schema".repeat(20),
+				parameters: Type.Object({ value: Type.String() }),
+			},
+		],
+		messages: [
+			{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: 1 },
+			{
+				role: "toolResult",
+				toolCallId: "call-1",
+				toolName: "read",
+				content: [{ type: "text", text: "tool output" }],
+				isError: false,
+				timestamp: 2,
+			},
+			{ role: "user", content: [{ type: "image", data: "abc123", mimeType: "image/png" }], timestamp: 3 },
+		],
+	};
+	const sourceMessages = [
+		...context.messages,
+		{
+			role: "custom",
+			customType: "gsd-memory",
+			content: "memory block",
+			display: false,
+			timestamp: 4,
+		} as AgentMessage,
+	];
+
+	const summary = buildTokenAuditSummary(context, sourceMessages);
+
+	assert.equal(summary.systemChars, "system prompt".length);
+	assert.equal(summary.toolCount, 2);
+	assert.equal(summary.messageCount, 3);
+	assert.equal(summary.toolResultChars, "tool output".length);
+	assert.equal(summary.imageCount, 1);
+	assert.ok(summary.toolSchemaChars > 0);
+	assert.ok(summary.customMessageChars > 0);
+	assert.ok(summary.estimatedInputTokens > 0);
+	assert.deepEqual(
+		summary.largestMessages.map((message) => Object.keys(message).sort()),
+		summary.largestMessages.map(() => ["chars", "index", "role", "type"]),
+	);
+	assert.equal(summary.largestTools[0].name, "large_tool");
+	assert.deepEqual(
+		summary.largestTools.map((tool) => Object.keys(tool).sort()),
+		summary.largestTools.map(() => ["chars", "name"]),
+	);
+	assert.deepEqual(summary.largestCustomMessages, [
+		{ index: 3, role: "custom", customType: "gsd-memory", chars: summary.largestCustomMessages[0].chars },
+	]);
+	assert.ok(!JSON.stringify(summary).includes("tool output"));
+	assert.ok(!JSON.stringify(summary).includes("memory block"));
+});
+
+test("maybeLogTokenAudit is opt-in and emits metadata only", () => {
+	const original = process.env.PI_TOKEN_AUDIT;
+	const originalWrite = process.stderr.write;
+	let written = "";
+	process.stderr.write = ((chunk: string | Uint8Array) => {
+		written += chunk.toString();
+		return true;
+	}) as typeof process.stderr.write;
+
+	try {
+		delete process.env.PI_TOKEN_AUDIT;
+		maybeLogTokenAudit({ messages: [{ role: "user", content: "secret prompt", timestamp: 1 }] }, []);
+		assert.equal(written, "");
+
+		process.env.PI_TOKEN_AUDIT = "1";
+		maybeLogTokenAudit({ systemPrompt: "hidden system", messages: [{ role: "user", content: "secret prompt", timestamp: 1 }] }, []);
+		assert.match(written, /"type":"token_audit"/);
+		assert.doesNotMatch(written, /secret prompt/);
+		assert.doesNotMatch(written, /hidden system/);
+	} finally {
+		process.stderr.write = originalWrite;
+		if (original === undefined) delete process.env.PI_TOKEN_AUDIT;
+		else process.env.PI_TOKEN_AUDIT = original;
+	}
+});
+
+test("provider payload audit summarizes post-hook payload without raw content", () => {
+	const payload = {
+		system: "secret system content",
+		tools: [{
+			type: "function",
+			function: {
+				name: "read",
+				description: "secret tool description",
+				parameters: { type: "object" },
+			},
+		}],
+		messages: [
+			{ role: "user", content: "secret user content" },
+			{ role: "assistant", content: [{ type: "text", text: "secret assistant content" }] },
+		],
+	};
+
+	const summary = buildProviderPayloadAuditSummary(payload);
+
+	assert.equal(summary.messageCount, 2);
+	assert.equal(summary.toolCount, 1);
+	assert.ok(summary.payloadChars > 0);
+	assert.ok(summary.toolSchemaChars > 0);
+	assert.deepEqual(summary.largestTools.map((tool) => tool.name), ["read"]);
+	assert.equal(JSON.stringify(summary).includes("secret"), false);
+});
+
+test("provider payload audit recognizes Gemini and Bedrock payload shapes", () => {
+	const gemini = buildProviderPayloadAuditSummary({
+		contents: [{ role: "user", parts: [{ text: "hidden gemini prompt" }] }],
+		config: {
+			tools: [{
+				functionDeclarations: [
+					{ name: "gsd_exec", description: "hidden declaration", parameters: { type: "object" } },
+				],
+			}],
+		},
+	});
+	const bedrock = buildProviderPayloadAuditSummary({
+		messages: [{ role: "user", content: [{ text: "hidden bedrock prompt" }] }],
+		toolConfig: {
+			tools: [
+				{ toolSpec: { name: "gsd_resume", description: "hidden tool", inputSchema: { json: {} } } },
+			],
+		},
+	});
+
+	assert.equal(gemini.messageCount, 1);
+	assert.equal(gemini.toolCount, 1);
+	assert.deepEqual(gemini.largestTools.map((tool) => tool.name), ["gsd_exec"]);
+	assert.equal(JSON.stringify(gemini).includes("hidden"), false);
+
+	assert.equal(bedrock.messageCount, 1);
+	assert.equal(bedrock.toolCount, 1);
+	assert.deepEqual(bedrock.largestTools.map((tool) => tool.name), ["gsd_resume"]);
+	assert.equal(JSON.stringify(bedrock).includes("hidden"), false);
+});
+
+test("provider payload audit logging is metadata-only", () => {
+	const original = process.env.PI_TOKEN_AUDIT;
+	const originalWrite = process.stderr.write;
+	let written = "";
+	process.env.PI_TOKEN_AUDIT = "1";
+	process.stderr.write = ((chunk: string | Uint8Array) => {
+		written += chunk.toString();
+		return true;
+	}) as typeof process.stderr.write;
+
+	try {
+		maybeLogProviderPayloadAudit({
+			messages: [{ role: "user", content: "raw prompt text must not log" }],
+			tools: [{ name: "bash", description: "raw tool description must not log" }],
+		}, "after");
+		assert.match(written, /"type":"token_audit_provider_payload"/);
+		assert.match(written, /"phase":"after"/);
+		assert.doesNotMatch(written, /raw prompt text/);
+		assert.doesNotMatch(written, /raw tool description/);
+	} finally {
+		process.stderr.write = originalWrite;
+		if (original === undefined) delete process.env.PI_TOKEN_AUDIT;
+		else process.env.PI_TOKEN_AUDIT = original;
+	}
+});

package/packages/pi-agent-core/src/token-audit.ts

@@ -0,0 +1,287 @@
+// Project/App: GSD-2
+// File Purpose: Provider-boundary token payload audit helpers.
+
+import type { Context, ImageContent, Message, TextContent, Tool } from "@gsd/pi-ai";
+import type { AgentMessage } from "./types.js";
+
+const CHARS_PER_TOKEN = 4;
+const LARGEST_MESSAGE_LIMIT = 5;
+const LARGEST_TOOL_LIMIT = 10;
+const LARGEST_CUSTOM_MESSAGE_LIMIT = 5;
+
+export interface TokenAuditMessageSummary {
+	index: number;
+	role: string;
+	type: string;
+	chars: number;
+}
+
+export interface TokenAuditToolSummary {
+	name: string;
+	chars: number;
+}
+
+export interface TokenAuditCustomMessageSummary {
+	index: number;
+	role: string;
+	customType?: string;
+	chars: number;
+}
+
+export interface TokenAuditSummary {
+	systemChars: number;
+	toolSchemaChars: number;
+	messageCharsByRole: Record<string, number>;
+	toolResultChars: number;
+	customMessageChars: number;
+	imageCount: number;
+	estimatedInputTokens: number;
+	messageCount: number;
+	toolCount: number;
+	largestMessages: TokenAuditMessageSummary[];
+	largestTools: TokenAuditToolSummary[];
+	largestCustomMessages: TokenAuditCustomMessageSummary[];
+}
+
+export interface ProviderPayloadAuditSummary {
+	payloadChars: number;
+	messageCharsByRole: Record<string, number>;
+	toolSchemaChars: number;
+	imageCount: number;
+	messageCount: number;
+	toolCount: number;
+	largestMessages: TokenAuditMessageSummary[];
+	largestTools: TokenAuditToolSummary[];
+}
+
+export function buildTokenAuditSummary(context: Context, sourceMessages: AgentMessage[]): TokenAuditSummary {
+	const systemChars = context.systemPrompt?.length ?? 0;
+	const toolSummaries = summarizeTools(context.tools ?? []);
+	const toolSchemaChars = toolSummaries.reduce((sum, tool) => sum + tool.chars, 0);
+	const messageSummaries = context.messages.map((message, index) => summarizeMessage(message, index));
+	const messageCharsByRole: Record<string, number> = {};
+	let toolResultChars = 0;
+	let imageCount = 0;
+
+	for (const summary of messageSummaries) {
+		messageCharsByRole[summary.role] = (messageCharsByRole[summary.role] ?? 0) + summary.chars;
+	}
+
+	for (const message of context.messages) {
+		imageCount += countImagesInMessage(message);
+		if (message.role === "toolResult") {
+			toolResultChars += countContentChars(message.content);
+		}
+	}
+
+	const customMessageSummaries = summarizeCustomMessages(sourceMessages);
+	const customMessageChars = customMessageSummaries.reduce((sum, message) => sum + message.chars, 0);
+
+	const totalChars =
+		systemChars +
+		toolSchemaChars +
+		messageSummaries.reduce((sum, message) => sum + message.chars, 0);
+
+	return {
+		systemChars,
+		toolSchemaChars,
+		messageCharsByRole,
+		toolResultChars,
+		customMessageChars,
+		imageCount,
+		estimatedInputTokens: Math.ceil(totalChars / CHARS_PER_TOKEN),
+		messageCount: context.messages.length,
+		toolCount: context.tools?.length ?? 0,
+		largestMessages: [...messageSummaries].sort((a, b) => b.chars - a.chars).slice(0, LARGEST_MESSAGE_LIMIT),
+		largestTools: [...toolSummaries].sort((a, b) => b.chars - a.chars).slice(0, LARGEST_TOOL_LIMIT),
+		largestCustomMessages: [...customMessageSummaries]
+			.sort((a, b) => b.chars - a.chars)
+			.slice(0, LARGEST_CUSTOM_MESSAGE_LIMIT),
+	};
+}
+
+export function maybeLogTokenAudit(context: Context, sourceMessages: AgentMessage[]): void {
+	if (process.env.PI_TOKEN_AUDIT !== "1") return;
+	const summary = buildTokenAuditSummary(context, sourceMessages);
+	process.stderr.write(`${JSON.stringify({ type: "token_audit", summary })}\n`);
+}
+
+export function buildProviderPayloadAuditSummary(payload: unknown): ProviderPayloadAuditSummary {
+	const record = asRecord(payload);
+	const messages = extractProviderMessages(record);
+	const tools = extractProviderTools(record);
+	const messageSummaries = messages.map((message, index) => summarizeProviderMessage(message, index));
+	const toolSummaries = tools.map((tool) => summarizeProviderTool(tool));
+	const messageCharsByRole: Record<string, number> = {};
+	let imageCount = 0;
+
+	for (const summary of messageSummaries) {
+		messageCharsByRole[summary.role] = (messageCharsByRole[summary.role] ?? 0) + summary.chars;
+	}
+	for (const message of messages) {
+		imageCount += countImagesInValue(message);
+	}
+
+	return {
+		payloadChars: safeJsonLength(payload),
+		messageCharsByRole,
+		toolSchemaChars: toolSummaries.reduce((sum, tool) => sum + tool.chars, 0),
+		imageCount,
+		messageCount: messages.length,
+		toolCount: tools.length,
+		largestMessages: [...messageSummaries].sort((a, b) => b.chars - a.chars).slice(0, LARGEST_MESSAGE_LIMIT),
+		largestTools: [...toolSummaries].sort((a, b) => b.chars - a.chars).slice(0, LARGEST_TOOL_LIMIT),
+	};
+}
+
+export function maybeLogProviderPayloadAudit(payload: unknown, phase: string): void {
+	if (process.env.PI_TOKEN_AUDIT !== "1") return;
+	const summary = buildProviderPayloadAuditSummary(payload);
+	process.stderr.write(`${JSON.stringify({ type: "token_audit_provider_payload", phase, summary })}\n`);
+}
+
+function summarizeMessage(message: Message, index: number): TokenAuditMessageSummary {
+	return {
+		index,
+		role: message.role,
+		type: messageType(message),
+		chars: safeJsonLength(message),
+	};
+}
+
+function messageType(message: Message): string {
+	if (message.role === "assistant") {
+		const types = new Set(message.content.map((content) => content.type));
+		return types.size > 0 ? Array.from(types).join("+") : "assistant";
+	}
+	if (message.role === "toolResult") return message.toolName || "toolResult";
+	if (typeof message.content === "string") return "text";
+	const types = new Set(message.content.map((content) => content.type));
+	return types.size > 0 ? Array.from(types).join("+") : "user";
+}
+
+function summarizeTools(tools: Tool[]): TokenAuditToolSummary[] {
+	return tools.map((tool) => ({
+		name: tool.name,
+		chars: safeJsonLength(tool),
+	}));
+}
+
+function summarizeCustomMessages(messages: AgentMessage[]): TokenAuditCustomMessageSummary[] {
+	return messages.flatMap((message, index) => {
+		if (!isCustomOrInjectedMessage(message)) return [];
+		const customType = (message as { customType?: unknown }).customType;
+		return [{
+			index,
+			role: String((message as { role?: unknown }).role ?? "unknown"),
+			customType: typeof customType === "string" ? customType : undefined,
+			chars: safeJsonLength(message),
+		}];
+	});
+}
+
+function summarizeProviderMessage(message: unknown, index: number): TokenAuditMessageSummary {
+	const record = asRecord(message);
+	const role = typeof record?.role === "string" ? record.role : "unknown";
+	const type = typeof record?.type === "string" ? record.type : role;
+	return { index, role, type, chars: safeJsonLength(message) };
+}
+
+function summarizeProviderTool(tool: unknown): TokenAuditToolSummary {
+	const record = asRecord(tool);
+	const name =
+		typeof record?.name === "string"
+			? record.name
+			: typeof asRecord(record?.function)?.name === "string"
+				? String(asRecord(record?.function)?.name)
+				: typeof asRecord(record?.toolSpec)?.name === "string"
+					? String(asRecord(record?.toolSpec)?.name)
+					: "unknown";
+	return { name, chars: safeJsonLength(tool) };
+}
+
+function extractProviderMessages(record: Record<string, unknown> | null): unknown[] {
+	const candidates = [
+		getArrayField(record, "messages"),
+		getArrayField(record, "input"),
+		getArrayField(record, "contents"),
+	].filter((value): value is unknown[] => value !== null);
+	if (candidates.length > 0) return candidates[0];
+	const input = record?.input;
+	return typeof input === "string" ? [{ role: "input", content: input }] : [];
+}
+
+function extractProviderTools(record: Record<string, unknown> | null): unknown[] {
+	const direct = getArrayField(record, "tools")
+		?? getNestedArrayField(record, ["config", "tools"])
+		?? getNestedArrayField(record, ["toolConfig", "tools"])
+		?? getNestedArrayField(record, ["tool_config", "tools"])
+		?? [];
+	return direct.flatMap(expandProviderToolEntry);
+}
+
+function expandProviderToolEntry(tool: unknown): unknown[] {
+	const record = asRecord(tool);
+	const functionDeclarations =
+		getArrayField(record, "functionDeclarations")
+		?? getArrayField(record, "function_declarations");
+	if (functionDeclarations) return functionDeclarations;
+	return [tool];
+}
+
+function countImagesInMessage(message: Message): number {
+	if (!("content" in message) || !Array.isArray(message.content)) return 0;
+	return message.content.filter((content): content is ImageContent => content.type === "image").length;
+}
+
+function countImagesInValue(value: unknown): number {
+	if (!value || typeof value !== "object") return 0;
+	if (Array.isArray(value)) {
+		return value.reduce((sum, item) => sum + countImagesInValue(item), 0);
+	}
+	const record = value as Record<string, unknown>;
+	let count = record.type === "image" || record.type === "input_image" ? 1 : 0;
+	for (const nested of Object.values(record)) {
+		count += countImagesInValue(nested);
+	}
+	return count;
+}
+
+function countContentChars(content: string | (TextContent | ImageContent)[]): number {
+	if (typeof content === "string") return content.length;
+	return content.reduce((sum, block) => {
+		if (block.type === "text") return sum + block.text.length;
+		if (block.type === "image") return sum + block.data.length;
+		return sum;
+	}, 0);
+}
+
+function isCustomOrInjectedMessage(message: AgentMessage): boolean {
+	const role = (message as { role?: unknown }).role;
+	return role === "custom" || role === "bashExecution" || role === "branchSummary" || role === "compactionSummary";
+}
+
+function asRecord(value: unknown): Record<string, unknown> | null {
+	return value && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : null;
+}
+
+function getArrayField(record: Record<string, unknown> | null, key: string): unknown[] | null {
+	const value = record?.[key];
+	return Array.isArray(value) ? value : null;
+}
+
+function getNestedArrayField(record: Record<string, unknown> | null, path: string[]): unknown[] | null {
+	let current: unknown = record;
+	for (const key of path) {
+		current = asRecord(current)?.[key];
+	}
+	return Array.isArray(current) ? current : null;
+}
+
+function safeJsonLength(value: unknown): number {
+	try {
+		return JSON.stringify(value)?.length ?? 0;
+	} catch {
+		return 0;
+	}
+}

package/packages/pi-agent-core/src/types.ts

@@ -134,6 +134,20 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 */
 	transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
 
+	/**
+	 * Optional final tool filter applied immediately before each provider call.
+	 *
+	 * Use this for runtime policy that depends on the fully active tool set.
+	 * The returned list is what token audit and the provider request both see.
+	 * Receives the post-transform AgentMessage context so policy can scope
+	 * request-local custom messages without inspecting provider payload text.
+	 */
+	filterTools?: (
+		tools: AgentTool<any>[],
+		signal?: AbortSignal,
+		messages?: AgentMessage[],
+	) => AgentTool<any>[] | Promise<AgentTool<any>[]>;
+
 	/**
 	 * Resolves an API key dynamically for each LLM call.
 	 *
@@ -320,23 +334,25 @@ export interface AgentContext {
 	tools?: AgentTool<any>[];
 }
 
+export type AgentAbortOrigin = "session-transition" | "user" | "timeout" | "unknown";
+
 /**
  * Events emitted by the Agent for UI updates.
  * These events provide fine-grained lifecycle information for messages, turns, and tool executions.
  */
 export type AgentEvent =
 	// Agent lifecycle
-	| { type: "agent_start" }
-	| { type: "agent_end"; messages: AgentMessage[] }
+	| { type: "agent_start"; sessionId?: string; turnId?: string }
+	| { type: "agent_end"; messages: AgentMessage[]; sessionId?: string; turnId?: string; abortOrigin?: AgentAbortOrigin }
 	// Turn lifecycle - a turn is one assistant response + any tool calls/results
-	| { type: "turn_start" }
-	| { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] }
+	| { type: "turn_start"; sessionId?: string; turnId?: string }
+	| { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[]; sessionId?: string; turnId?: string }
 	// Message lifecycle - emitted for user, assistant, and toolResult messages
-	| { type: "message_start"; message: AgentMessage }
+	| { type: "message_start"; message: AgentMessage; sessionId?: string; turnId?: string }
 	// Only emitted for assistant messages during streaming
-	| { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }
-	| { type: "message_end"; message: AgentMessage }
+	| { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent; sessionId?: string; turnId?: string }
+	| { type: "message_end"; message: AgentMessage; sessionId?: string; turnId?: string }
 	// Tool execution lifecycle
-	| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
-	| { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any }
-	| { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean };
+	| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any; sessionId?: string; turnId?: string }
+	| { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any; sessionId?: string; turnId?: string }
+	| { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean; sessionId?: string; turnId?: string };