fix(coding-agent): add explicit initiator attribution for Copilot headers (#246)

* fix(coding-agent): add explicit initiator attribution

Use message-level attribution for Copilot X-Initiator with role-based fallback, and persist attribution across custom/hook session paths.

Fixes #237

* fix(coding-agent): preserve legacy custom attribution fallback

* fix(coding-agent): remove async-result role special-case

* fix(coding-agent): inherit before_agent_start attribution from prompt

* test(coding-agent): tighten typing in attribution regressions

packages/ai/src/providers/github-copilot-headers.ts

@@ -8,6 +8,14 @@ export function inferCopilotInitiator(messages: unknown[]): "user" | "agent" {
 	if (messages.length === 0) return "user";
 
 	const last = messages[messages.length - 1] as Record<string, unknown>;
+	const attribution = last.attribution;
+	if (typeof attribution === "string") {
+		const normalizedAttribution = attribution.trim().toLowerCase();
+		if (normalizedAttribution === "user" || normalizedAttribution === "agent") {
+			return normalizedAttribution;
+		}
+	}
+
 	const role = last.role as string | undefined;
 	if (!role) return "user";
 

packages/ai/src/types.ts

@@ -112,6 +112,8 @@ export type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh";
 /** Token budgets for each thinking level (token-based providers only) */
 export type ThinkingBudgets = { [key in ThinkingLevel]?: number };
 
+export type MessageAttribution = "user" | "agent";
+
 export type ToolChoice =
 	| "auto"
 	| "none"
@@ -257,12 +259,16 @@ export interface UserMessage {
 	content: string | (TextContent | ImageContent)[];
 	/** True if the message was injected by the system (e.g., auto-continue). */
 	synthetic?: boolean;
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 	timestamp: number; // Unix timestamp in milliseconds
 }
 
 export interface DeveloperMessage {
 	role: "developer";
 	content: string | (TextContent | ImageContent)[];
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 	timestamp: number; // Unix timestamp in milliseconds
 }
 
@@ -287,6 +293,8 @@ export interface ToolResultMessage<TDetails = any> {
 	content: (TextContent | ImageContent)[]; // Supports text and images
 	details?: TDetails;
 	isError: boolean;
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 	/** Timestamp when output was pruned (ms since epoch). Undefined if unpruned. */
 	prunedAt?: number;
 	timestamp: number; // Unix timestamp in milliseconds

packages/ai/test/github-copilot-headers.test.ts

@@ -66,6 +66,19 @@ describe("inferCopilotInitiator", () => {
 		expect(inferCopilotInitiator(messages)).toBe("user");
 	});
 
+	it("prefers explicit attribution over role when attribution is agent", () => {
+		const messages: Message[] = [
+			{ role: "user", content: "internal reminder", attribution: "agent", timestamp: Date.now() },
+		];
+		expect(inferCopilotInitiator(messages)).toBe("agent");
+	});
+
+	it("prefers explicit attribution over role when attribution is user", () => {
+		const messages: Message[] = [
+			{ role: "developer", content: "forward user note", attribution: "user", timestamp: Date.now() },
+		];
+		expect(inferCopilotInitiator(messages)).toBe("user");
+	});
 	it("returns 'agent' when last message is user but last content block is tool_result", () => {
 		const messages: unknown[] = [
 			{

packages/coding-agent/src/extensibility/extensions/loader.ts

@@ -174,7 +174,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime {
 	}
 
 	sendMessage<T = unknown>(
-		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
+		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 		options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" },
 	): void {
 		this.runtime.sendMessage(message, options);

packages/coding-agent/src/extensibility/extensions/types.ts

@@ -824,7 +824,7 @@ export interface ToolResultEventResult {
 }
 
 export interface BeforeAgentStartEventResult {
-	message?: Pick<CustomMessage, "customType" | "content" | "display" | "details">;
+	message?: Pick<CustomMessage, "customType" | "content" | "display" | "details" | "attribution">;
 	/** Replace the system prompt for this turn. If multiple extensions return this, they are chained. */
 	systemPrompt?: string;
 }
@@ -1015,7 +1015,7 @@ export interface ExtensionAPI {
 
 	/** Send a custom message to the session. */
 	sendMessage<T = unknown>(
-		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
+		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 		options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" },
 	): void;
 
@@ -1184,7 +1184,7 @@ export interface ExtensionShortcut {
 type HandlerFn = (...args: unknown[]) => Promise<unknown>;
 
 export type SendMessageHandler = <T = unknown>(
-	message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
+	message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 	options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" },
 ) => void;
 

packages/coding-agent/src/extensibility/hooks/loader.ts

@@ -23,7 +23,7 @@ type HandlerFn = (...args: unknown[]) => Promise<unknown>;
  * Send message handler type for pi.sendMessage().
  */
 export type SendMessageHandler = <T = unknown>(
-	message: Pick<HookMessage<T>, "customType" | "content" | "display" | "details">,
+	message: Pick<HookMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 	options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" },
 ) => void;
 

packages/coding-agent/src/extensibility/hooks/types.ts

@@ -585,7 +585,7 @@ export interface ToolResultEventResult {
  */
 export interface BeforeAgentStartEventResult {
 	/** Message to inject into context (persisted to session, visible in TUI) */
-	message?: Pick<HookMessage, "customType" | "content" | "display" | "details">;
+	message?: Pick<HookMessage, "customType" | "content" | "display" | "details" | "attribution">;
 }
 
 /** Return type for session_before_switch handlers */
@@ -733,12 +733,13 @@ export interface HookAPI {
 	 * @param message.content - Message content (string or TextContent/ImageContent array)
 	 * @param message.display - Whether to show in TUI (true = styled display, false = hidden)
 	 * @param message.details - Optional hook-specific metadata (not sent to LLM)
+	 * @param message.attribution - Who initiated the message for billing/attribution semantics ("user" | "agent")
 	 * @param options.triggerTurn - If true and agent is idle, triggers a new LLM turn. Default: false.
 	 *                              If agent is streaming, message is queued and triggerTurn is ignored.
 	 * @param options.deliverAs - How to deliver the message: "steer" or "followUp".
 	 */
 	sendMessage<T = unknown>(
-		message: Pick<HookMessage<T>, "customType" | "content" | "display" | "details">,
+		message: Pick<HookMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 		options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" },
 	): void;
 

packages/coding-agent/src/modes/controllers/input-controller.ts

@@ -225,6 +225,7 @@ export class InputController {
 								content: message,
 								display: true,
 								details,
+								attribution: "user",
 							},
 							{ streamingBehavior: "followUp" },
 						);

packages/coding-agent/src/sdk.ts

@@ -763,6 +763,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 							customType: "async-result",
 							content: message,
 							display: true,
+							attribution: "agent",
 							details: {
 								jobId,
 								type: job?.type,

packages/coding-agent/src/session/agent-session.ts

@@ -422,6 +422,7 @@ export class AgentSession {
 				content: reminderText,
 				display: false,
 				details: { toolName: action.sourceToolName },
+				attribution: "agent",
 				timestamp: Date.now(),
 			});
 		});
@@ -597,6 +598,7 @@ export class AgentSession {
 										content: injection.content,
 										display: false,
 										details,
+										attribution: "agent",
 										timestamp: Date.now(),
 									});
 									this.sessionManager.appendCustomMessageEntry(
@@ -604,6 +606,7 @@ export class AgentSession {
 										injection.content,
 										false,
 										details,
+										"agent",
 									);
 									this.#markTtsrInjected(details.rules);
 								}
@@ -642,6 +645,7 @@ export class AgentSession {
 					event.message.content,
 					event.message.display,
 					event.message.details,
+					event.message.attribution ?? "agent",
 				);
 				if (event.message.role === "custom" && event.message.customType === "ttsr-injection") {
 					this.#markTtsrInjected(this.#extractTtsrRuleNames(event.message.details));
@@ -1011,6 +1015,7 @@ export class AgentSession {
 			content: injection.content,
 			display: false,
 			details: { rules: injection.rules.map(rule => rule.name) },
+			attribution: "agent",
 			timestamp: Date.now(),
 		});
 		this.#ensureTtsrResumePromise();
@@ -1809,6 +1814,7 @@ export class AgentSession {
 			customType: "plan-mode-reference",
 			content,
 			display: false,
+			attribution: "agent",
 			timestamp: Date.now(),
 		};
 	}
@@ -1849,6 +1855,7 @@ export class AgentSession {
 			customType: "plan-mode-context",
 			content,
 			display: false,
+			attribution: "agent",
 			timestamp: Date.now(),
 		};
 	}
@@ -1910,8 +1917,8 @@ export class AgentSession {
 		}
 
 		const message = options?.synthetic
-			? { role: "developer" as const, content: userContent, timestamp: Date.now() }
-			: { role: "user" as const, content: userContent, timestamp: Date.now() };
+			? { role: "developer" as const, content: userContent, attribution: "agent" as const, timestamp: Date.now() }
+			: { role: "user" as const, content: userContent, attribution: "user" as const, timestamp: Date.now() };
 
 		await this.#promptWithMessage(message, expandedText, options);
 		if (!options?.synthetic) {
@@ -1920,7 +1927,7 @@ export class AgentSession {
 	}
 
 	async promptCustomMessage<T = unknown>(
-		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
+		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 		options?: Pick<PromptOptions, "streamingBehavior" | "toolChoice">,
 	): Promise<void> {
 		const textContent =
@@ -1945,6 +1952,7 @@ export class AgentSession {
 			content: message.content,
 			display: message.display,
 			details: message.details,
+			attribution: message.attribution ?? "agent",
 			timestamp: Date.now(),
 		};
 
@@ -2033,6 +2041,8 @@ export class AgentSession {
 					this.#baseSystemPrompt,
 				);
 				if (result?.messages) {
+					const promptAttribution: "user" | "agent" | undefined =
+						"attribution" in message ? message.attribution : undefined;
 					for (const msg of result.messages) {
 						messages.push({
 							role: "custom",
@@ -2040,6 +2050,7 @@ export class AgentSession {
 							content: msg.content,
 							display: msg.display,
 							details: msg.details,
+							attribution: msg.attribution ?? promptAttribution ?? (message.role === "user" ? "user" : "agent"),
 							timestamp: Date.now(),
 						});
 					}
@@ -2239,6 +2250,7 @@ export class AgentSession {
 		this.agent.steer({
 			role: "user",
 			content,
+			attribution: "user",
 			timestamp: Date.now(),
 		});
 	}
@@ -2256,6 +2268,7 @@ export class AgentSession {
 		this.agent.followUp({
 			role: "user",
 			content,
+			attribution: "user",
 			timestamp: Date.now(),
 		});
 	}
@@ -2286,7 +2299,7 @@ export class AgentSession {
 	 * - Not streaming + no trigger: appends to state/session, no turn
 	 */
 	async sendCustomMessage<T = unknown>(
-		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details">,
+		message: Pick<CustomMessage<T>, "customType" | "content" | "display" | "details" | "attribution">,
 		options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" },
 	): Promise<void> {
 		const appMessage: CustomMessage<T> = {
@@ -2295,6 +2308,7 @@ export class AgentSession {
 			content: message.content,
 			display: message.display,
 			details: message.details,
+			attribution: message.attribution ?? "agent",
 			timestamp: Date.now(),
 		};
 		if (this.isStreaming) {
@@ -2322,6 +2336,7 @@ export class AgentSession {
 			message.content,
 			message.display,
 			message.details,
+			message.attribution ?? "agent",
 		);
 	}
 
@@ -3212,7 +3227,7 @@ Be thorough - include exact file paths, function names, error messages, and tech
 
 			// Inject the handoff document as a custom message
 			const handoffContent = `<handoff-context>\n${handoffText}\n</handoff-context>\n\nThe above is a handoff document from a previous session. Use this context to continue the work seamlessly.`;
-			this.sessionManager.appendCustomMessageEntry("handoff", handoffContent, true);
+			this.sessionManager.appendCustomMessageEntry("handoff", handoffContent, true, undefined, "agent");
 
 			// Rebuild agent messages from session
 			const sessionContext = this.sessionManager.buildSessionContext();
@@ -3309,6 +3324,7 @@ Be thorough - include exact file paths, function names, error messages, and tech
 		this.agent.appendMessage({
 			role: "developer",
 			content: [{ type: "text", text: reminder }],
+			attribution: "agent",
 			timestamp: Date.now(),
 		});
 		this.#scheduleAgentContinue({ generation: this.#promptGeneration });
@@ -3339,9 +3355,10 @@ Be thorough - include exact file paths, function names, error messages, and tech
 			content: report,
 			display: false,
 			details,
+			attribution: "agent",
 			timestamp: Date.now(),
 		});
-		this.sessionManager.appendCustomMessageEntry("rewind-report", report, false, details);
+		this.sessionManager.appendCustomMessageEntry("rewind-report", report, false, details, "agent");
 		this.#checkpointState = undefined;
 		this.#pendingRewindReport = undefined;
 	}
@@ -3460,6 +3477,7 @@ Be thorough - include exact file paths, function names, error messages, and tech
 		this.agent.appendMessage({
 			role: "developer",
 			content: [{ type: "text", text: reminder }],
+			attribution: "agent",
 			timestamp: Date.now(),
 		});
 		this.#scheduleAgentContinue({ generation: this.#promptGeneration });
@@ -3876,6 +3894,7 @@ Be thorough - include exact file paths, function names, error messages, and tech
 					{
 						role: "developer",
 						content: [{ type: "text", text: "Continue if you have next steps." }],
+						attribution: "agent",
 						timestamp: Date.now(),
 					},
 					"Continue if you have next steps.",

packages/coding-agent/src/session/compaction/branch-summarization.ts

@@ -149,7 +149,14 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
 			return entry.message;
 
 		case "custom_message":
-			return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
+			return createCustomMessage(
+				entry.customType,
+				entry.content,
+				entry.display,
+				entry.details,
+				entry.timestamp,
+				entry.attribution,
+			);
 
 		case "branch_summary":
 			return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);

packages/coding-agent/src/session/compaction/compaction.ts

@@ -81,7 +81,14 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
 		return entry.message;
 	}
 	if (entry.type === "custom_message") {
-		return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
+		return createCustomMessage(
+			entry.customType,
+			entry.content,
+			entry.display,
+			entry.details,
+			entry.timestamp,
+			entry.attribution,
+		);
 	}
 	if (entry.type === "branch_summary") {
 		return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);

packages/coding-agent/src/session/messages.ts

@@ -5,7 +5,7 @@
  * and provides a transformer to convert them to LLM-compatible messages.
  */
 import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
-import type { ImageContent, Message, TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
+import type { ImageContent, Message, MessageAttribution, TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
 import { renderPromptTemplate } from "../config/prompt-templates";
 import branchSummaryContextPrompt from "../prompts/compaction/branch-summary-context.md" with { type: "text" };
 import compactionSummaryContextPrompt from "../prompts/compaction/compaction-summary-context.md" with { type: "text" };
@@ -75,6 +75,8 @@ export interface CustomMessage<T = unknown> {
 	content: string | (TextContent | ImageContent)[];
 	display: boolean;
 	details?: T;
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 	timestamp: number;
 }
 
@@ -87,6 +89,8 @@ export interface HookMessage<T = unknown> {
 	content: string | (TextContent | ImageContent)[];
 	display: boolean;
 	details?: T;
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 	timestamp: number;
 }
 
@@ -206,6 +210,7 @@ export function createCustomMessage(
 	display: boolean,
 	details: unknown | undefined,
 	timestamp: string,
+	attribution?: MessageAttribution,
 ): CustomMessage {
 	return {
 		role: "custom",
@@ -213,6 +218,7 @@ export function createCustomMessage(
 		content,
 		display,
 		details,
+		attribution,
 		timestamp: new Date(timestamp).getTime(),
 	};
 }
@@ -236,6 +242,7 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 					return {
 						role: "user",
 						content: [{ type: "text", text: bashExecutionToText(m) }],
+						attribution: "user",
 						timestamp: m.timestamp,
 					};
 				case "pythonExecution":
@@ -245,14 +252,18 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 					return {
 						role: "user",
 						content: [{ type: "text", text: pythonExecutionToText(m) }],
+						attribution: "user",
 						timestamp: m.timestamp,
 					};
 				case "custom":
 				case "hookMessage": {
 					const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;
+					const role = "user";
+					const attribution = m.attribution;
 					return {
-						role: "user",
+						role,
 						content,
+						attribution,
 						timestamp: m.timestamp,
 					};
 				}
@@ -265,6 +276,7 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 								text: renderPromptTemplate(BRANCH_SUMMARY_TEMPLATE, { summary: m.summary }),
 							},
 						],
+						attribution: "agent",
 						timestamp: m.timestamp,
 					};
 				case "compactionSummary":
@@ -276,6 +288,7 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 								text: renderPromptTemplate(COMPACTION_SUMMARY_TEMPLATE, { summary: m.summary }),
 							},
 						],
+						attribution: "agent",
 						timestamp: m.timestamp,
 					};
 				case "fileMention": {
@@ -296,17 +309,21 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 					return {
 						role: "user",
 						content,
+						attribution: "user",
 						timestamp: m.timestamp,
 					};
 				}
 				case "user":
+					return { ...m, attribution: m.attribution ?? "user" };
 				case "developer":
+					return { ...m, attribution: m.attribution ?? "agent" };
 				case "assistant":
 					return m;
 				case "toolResult":
 					return {
 						...m,
 						content: getPrunedToolResultContent(m as ToolResultMessage),
+						attribution: m.attribution ?? "agent",
 					};
 				default:
 					// biome-ignore lint/correctness/noSwitchDeclarations: fine

packages/coding-agent/src/session/session-manager.ts

@@ -2,7 +2,7 @@ import * as fs from "node:fs";
 import * as os from "node:os";
 import * as path from "node:path";
 import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
-import type { ImageContent, Message, TextContent, Usage } from "@oh-my-pi/pi-ai";
+import type { ImageContent, Message, MessageAttribution, TextContent, Usage } from "@oh-my-pi/pi-ai";
 import { getTerminalId } from "@oh-my-pi/pi-tui";
 import {
 	getBlobsDir,
@@ -151,7 +151,7 @@ export interface ModeChangeEntry extends SessionEntryBase {
  * Use customType to identify your extension's entries.
  *
  * Unlike CustomEntry, this DOES participate in LLM context.
- * The content is converted to a user message in buildSessionContext().
+ * The content participates in LLM context through convertToLlm().
  * Use details for extension-specific metadata (not sent to LLM).
  *
  * display controls TUI rendering:
@@ -164,6 +164,8 @@ export interface CustomMessageEntry<T = unknown> extends SessionEntryBase {
 	content: string | (TextContent | ImageContent)[];
 	details?: T;
 	display: boolean;
+	/** Who initiated this message for billing/attribution semantics. */
+	attribution?: MessageAttribution;
 }
 
 /** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
@@ -483,7 +485,14 @@ export function buildSessionContext(
 			messages.push(entry.message);
 		} else if (entry.type === "custom_message") {
 			messages.push(
-				createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
+				createCustomMessage(
+					entry.customType,
+					entry.content,
+					entry.display,
+					entry.details,
+					entry.timestamp,
+					entry.attribution,
+				),
 			);
 		} else if (entry.type === "branch_summary" && entry.summary) {
 			messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
@@ -1825,6 +1834,7 @@ export class SessionManager {
 	 * @param content Message content (string or TextContent/ImageContent array)
 	 * @param display Whether to show in TUI (true = styled display, false = hidden)
 	 * @param details Optional extension-specific metadata (not sent to LLM)
+	 * @param attribution Who initiated this message for billing/attribution semantics
 	 * @returns Entry id
 	 */
 	appendCustomMessageEntry<T = unknown>(
@@ -1832,6 +1842,7 @@ export class SessionManager {
 		content: string | (TextContent | ImageContent)[],
 		display: boolean,
 		details?: T,
+		attribution: MessageAttribution = "agent",
 	): string {
 		const entry: CustomMessageEntry<T> = {
 			type: "custom_message",
@@ -1839,6 +1850,7 @@ export class SessionManager {
 			content,
 			display,
 			details,
+			attribution,
 			id: generateId(this.#byId),
 			parentId: this.#leafId,
 			timestamp: new Date().toISOString(),

packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts

@@ -0,0 +1,161 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
+import * as path from "node:path";
+import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core";
+import { type AssistantMessage, getBundledModel, type Message } from "@oh-my-pi/pi-ai";
+import { inferCopilotInitiator } from "@oh-my-pi/pi-ai/providers/github-copilot-headers";
+import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
+import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
+import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
+import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions";
+import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
+import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
+import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
+import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
+import { TempDir } from "@oh-my-pi/pi-utils";
+
+class MockAssistantStream extends AssistantMessageEventStream {}
+
+function createAssistantMessage(text: string): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [{ type: "text", text }],
+		api: "anthropic-messages",
+		provider: "anthropic",
+		model: "mock",
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: "stop",
+		timestamp: Date.now(),
+	};
+}
+
+describe("AgentSession before_agent_start attribution fallback", () => {
+	let tempDir: TempDir;
+	let session: AgentSession;
+	let modelRegistry: ModelRegistry;
+
+	const injectedText = "before-agent-start injected message";
+
+	beforeEach(async () => {
+		tempDir = TempDir.createSync("@pi-before-agent-start-attribution-");
+		const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db"));
+		authStorage.setRuntimeApiKey("anthropic", "test-key");
+		modelRegistry = new ModelRegistry(authStorage);
+	});
+
+	afterEach(async () => {
+		vi.restoreAllMocks();
+		if (session) {
+			await session.dispose();
+		}
+		tempDir.removeSync();
+	});
+
+	function createSession() {
+		const emitBeforeAgentStart = vi.fn().mockResolvedValue({
+			messages: [
+				{
+					customType: "before-start",
+					content: injectedText,
+					display: false,
+				},
+			],
+		});
+		const extensionRunner = {
+			emitBeforeAgentStart,
+			emit: vi.fn().mockResolvedValue(undefined),
+		} as unknown as ExtensionRunner;
+
+		const model = getBundledModel("anthropic", "claude-sonnet-4-5");
+		if (!model) throw new Error("Expected claude-sonnet-4-5 model to exist");
+
+		const agent = new Agent({
+			getApiKey: () => "test-key",
+			initialState: {
+				model,
+				systemPrompt: "Test",
+				tools: [],
+				messages: [],
+			},
+			streamFn: () => {
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					stream.push({ type: "start", partial: createAssistantMessage("") });
+					stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") });
+				});
+				return stream;
+			},
+		});
+
+		session = new AgentSession({
+			agent,
+			sessionManager: SessionManager.inMemory(),
+			settings: Settings.isolated({ "compaction.enabled": false }),
+			modelRegistry,
+			extensionRunner,
+		});
+
+		return { emitBeforeAgentStart };
+	}
+
+	function findBeforeStartInjection(messages: AgentMessage[]): AgentMessage | undefined {
+		return messages.find(message => message.role === "custom" && message.customType === "before-start");
+	}
+
+	function findBeforeStartInjectionLlm(messages: Message[]): Message | undefined {
+		return messages.find(message => {
+			if (message.role === "assistant") return false;
+			if (typeof message.content === "string") return message.content === injectedText;
+			return message.content.some(block => block.type === "text" && block.text === injectedText);
+		});
+	}
+	it("defaults before_agent_start message attribution to user for user prompts", async () => {
+		const { emitBeforeAgentStart } = createSession();
+
+		await session.prompt("hello from user");
+
+		expect(emitBeforeAgentStart).toHaveBeenCalledTimes(1);
+		const injectedMessage = findBeforeStartInjection(session.messages);
+		expect(injectedMessage).toBeDefined();
+		if (!injectedMessage || injectedMessage.role !== "custom") {
+			throw new Error("Expected injected custom message in session state");
+		}
+
+		const llmMessages = convertToLlm(session.messages.filter(message => message.role !== "assistant"));
+		const llmInjected = findBeforeStartInjectionLlm(llmMessages);
+		expect(llmInjected).toBeDefined();
+		if (!llmInjected || llmInjected.role === "assistant") {
+			throw new Error("Expected injected message in converted LLM context");
+		}
+		expect(llmInjected.attribution).toBe("user");
+		expect(inferCopilotInitiator(llmMessages)).toBe("user");
+	});
+
+	it("defaults before_agent_start message attribution to agent for synthetic prompts", async () => {
+		const { emitBeforeAgentStart } = createSession();
+
+		await session.prompt("internal reminder", { synthetic: true });
+
+		expect(emitBeforeAgentStart).toHaveBeenCalledTimes(1);
+		const injectedMessage = findBeforeStartInjection(session.messages);
+		expect(injectedMessage).toBeDefined();
+		if (!injectedMessage || injectedMessage.role !== "custom") {
+			throw new Error("Expected injected custom message in session state");
+		}
+
+		const llmMessages = convertToLlm(session.messages.filter(message => message.role !== "assistant"));
+		const llmInjected = findBeforeStartInjectionLlm(llmMessages);
+		expect(llmInjected).toBeDefined();
+		if (!llmInjected || llmInjected.role === "assistant") {
+			throw new Error("Expected injected message in converted LLM context");
+		}
+		expect(llmInjected.attribution).toBe("agent");
+		expect(inferCopilotInitiator(llmMessages)).toBe("agent");
+	});
+});

packages/coding-agent/test/session-manager/build-context.test.ts

@@ -76,6 +76,45 @@ describe("buildSessionContext", () => {
 			expect(ctx.messages[0].role).toBe("user");
 		});
 
+		it("rehydrates custom_message attribution from entries", () => {
+			const entries: SessionEntry[] = [
+				{
+					type: "custom_message",
+					id: "1",
+					parentId: null,
+					timestamp: "2025-01-01T00:00:00Z",
+					customType: "skill-prompt",
+					content: "Summarize this file",
+					display: true,
+					attribution: "user",
+				},
+			];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.messages).toHaveLength(1);
+			const customMessage = ctx.messages[0];
+			expect(customMessage?.role).toBe("custom");
+			if (customMessage?.role !== "custom") throw new Error("Expected custom message");
+			expect(customMessage.attribution).toBe("user");
+		});
+		it("preserves missing custom_message attribution on rehydration", () => {
+			const entries: SessionEntry[] = [
+				{
+					type: "custom_message",
+					id: "1",
+					parentId: null,
+					timestamp: "2025-01-01T00:00:00Z",
+					customType: "skill-prompt",
+					content: "Summarize this file",
+					display: true,
+				},
+			];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.messages).toHaveLength(1);
+			const customMessage = ctx.messages[0];
+			expect(customMessage?.role).toBe("custom");
+			if (customMessage?.role !== "custom") throw new Error("Expected custom message");
+			expect(customMessage.attribution).toBeUndefined();
+		});
 		it("simple conversation", () => {
 			const entries: SessionEntry[] = [
 				msg("1", null, "user", "hello"),

packages/coding-agent/test/session-messages.test.ts

@@ -0,0 +1,95 @@
+import { describe, expect, it } from "bun:test";
+import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
+import type { Message } from "@oh-my-pi/pi-ai";
+import { inferCopilotInitiator } from "@oh-my-pi/pi-ai/providers/github-copilot-headers";
+import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
+
+function expectAttribution(message: Message | undefined, expected: "user" | "agent" | undefined): void {
+	expect(message).toBeDefined();
+	if (!message) return;
+	if (message.role === "assistant") {
+		throw new Error("Assistant messages do not expose attribution");
+	}
+	expect(message.attribution).toBe(expected);
+}
+
+describe("convertToLlm custom message mapping", () => {
+	it("uses async-result attribution without special role mapping", () => {
+		const messages: AgentMessage[] = [
+			{
+				role: "custom",
+				customType: "async-result",
+				content: "Background task completed",
+				display: true,
+				attribution: "agent",
+				timestamp: Date.now(),
+			},
+		];
+
+		const converted = convertToLlm(messages);
+
+		expect(converted).toHaveLength(1);
+		expect(converted[0]?.role).toBe("user");
+		expectAttribution(converted[0], "agent");
+		expect(inferCopilotInitiator(converted)).toBe("agent");
+	});
+
+	it("preserves missing attribution for legacy custom messages", () => {
+		const messages: AgentMessage[] = [
+			{
+				role: "custom",
+				customType: "skill-prompt",
+				content: "Run this skill with my arguments",
+				display: true,
+				timestamp: Date.now(),
+			},
+		];
+
+		const converted = convertToLlm(messages);
+
+		expect(converted).toHaveLength(1);
+		expect(converted[0]?.role).toBe("user");
+		expectAttribution(converted[0], undefined);
+		expect(inferCopilotInitiator(converted)).toBe("user");
+	});
+
+	it("uses explicit agent attribution for custom messages", () => {
+		const messages: AgentMessage[] = [
+			{
+				role: "custom",
+				customType: "ttsr-injection",
+				content: "<system-reminder>Read file</system-reminder>",
+				display: false,
+				attribution: "agent",
+				timestamp: Date.now(),
+			},
+		];
+
+		const converted = convertToLlm(messages);
+
+		expect(converted).toHaveLength(1);
+		expect(converted[0]?.role).toBe("user");
+		expectAttribution(converted[0], "agent");
+		expect(inferCopilotInitiator(converted)).toBe("agent");
+	});
+
+	it("allows custom messages to opt into user attribution", () => {
+		const messages: AgentMessage[] = [
+			{
+				role: "custom",
+				customType: "skill-prompt",
+				content: "Run this skill with my arguments",
+				display: true,
+				attribution: "user",
+				timestamp: Date.now(),
+			},
+		];
+
+		const converted = convertToLlm(messages);
+
+		expect(converted).toHaveLength(1);
+		expect(converted[0]?.role).toBe("user");
+		expectAttribution(converted[0], "user");
+		expect(inferCopilotInitiator(converted)).toBe("user");
+	});
+});