From 3119bfced5c867713f53575df6af99b73e86f660 Mon Sep 17 00:00:00 2001 From: maximhar Date: Mon, 2 Mar 2026 18:01:44 +0200 Subject: [PATCH] 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 --- .../src/providers/github-copilot-headers.ts | 8 + packages/ai/src/types.ts | 8 + .../ai/test/github-copilot-headers.test.ts | 13 ++ .../src/extensibility/extensions/loader.ts | 2 +- .../src/extensibility/extensions/types.ts | 6 +- .../src/extensibility/hooks/loader.ts | 2 +- .../src/extensibility/hooks/types.ts | 5 +- .../src/modes/controllers/input-controller.ts | 1 + packages/coding-agent/src/sdk.ts | 1 + .../coding-agent/src/session/agent-session.ts | 31 +++- .../compaction/branch-summarization.ts | 9 +- .../src/session/compaction/compaction.ts | 9 +- packages/coding-agent/src/session/messages.ts | 21 ++- .../src/session/session-manager.ts | 18 +- ...ion-before-agent-start-attribution.test.ts | 161 ++++++++++++++++++ .../session-manager/build-context.test.ts | 39 +++++ .../test/session-messages.test.ts | 95 +++++++++++ 17 files changed, 409 insertions(+), 20 deletions(-) create mode 100644 packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts create mode 100644 packages/coding-agent/test/session-messages.test.ts diff --git a/packages/ai/src/providers/github-copilot-headers.ts b/packages/ai/src/providers/github-copilot-headers.ts index bddd59a60..9a3c0d3fc 100644 --- a/packages/ai/src/providers/github-copilot-headers.ts +++ b/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; + 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"; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 05ff67212..feb9eb786 100644 --- a/packages/ai/src/types.ts +++ b/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 { 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 diff --git a/packages/ai/test/github-copilot-headers.test.ts b/packages/ai/test/github-copilot-headers.test.ts index ad894593d..cca75457a 100644 --- a/packages/ai/test/github-copilot-headers.test.ts +++ b/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[] = [ { diff --git a/packages/coding-agent/src/extensibility/extensions/loader.ts b/packages/coding-agent/src/extensibility/extensions/loader.ts index 8584b00d1..a525d821e 100644 --- a/packages/coding-agent/src/extensibility/extensions/loader.ts +++ b/packages/coding-agent/src/extensibility/extensions/loader.ts @@ -174,7 +174,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { } sendMessage( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ): void { this.runtime.sendMessage(message, options); diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 6431d71b6..389b4c986 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -824,7 +824,7 @@ export interface ToolResultEventResult { } export interface BeforeAgentStartEventResult { - message?: Pick; + message?: Pick; /** 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( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "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; export type SendMessageHandler = ( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ) => void; diff --git a/packages/coding-agent/src/extensibility/hooks/loader.ts b/packages/coding-agent/src/extensibility/hooks/loader.ts index e6c7f88b7..e876404b0 100644 --- a/packages/coding-agent/src/extensibility/hooks/loader.ts +++ b/packages/coding-agent/src/extensibility/hooks/loader.ts @@ -23,7 +23,7 @@ type HandlerFn = (...args: unknown[]) => Promise; * Send message handler type for pi.sendMessage(). */ export type SendMessageHandler = ( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" }, ) => void; diff --git a/packages/coding-agent/src/extensibility/hooks/types.ts b/packages/coding-agent/src/extensibility/hooks/types.ts index 2270c1cd0..9fffefe14 100644 --- a/packages/coding-agent/src/extensibility/hooks/types.ts +++ b/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; + message?: Pick; } /** 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( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" }, ): void; diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index ee987979f..717e2f119 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/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" }, ); diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 22b30d535..b167f3bbd 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/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, diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 00e9fbfa8..e95394c4b 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/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( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: Pick, ): Promise { 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( - message: Pick, "customType" | "content" | "display" | "details">, + message: Pick, "customType" | "content" | "display" | "details" | "attribution">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn" }, ): Promise { const appMessage: CustomMessage = { @@ -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 = `\n${handoffText}\n\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.", diff --git a/packages/coding-agent/src/session/compaction/branch-summarization.ts b/packages/coding-agent/src/session/compaction/branch-summarization.ts index 1219b2271..6bcf86485 100644 --- a/packages/coding-agent/src/session/compaction/branch-summarization.ts +++ b/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); diff --git a/packages/coding-agent/src/session/compaction/compaction.ts b/packages/coding-agent/src/session/compaction/compaction.ts index 25ca3edd7..09ba2ce57 100644 --- a/packages/coding-agent/src/session/compaction/compaction.ts +++ b/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); diff --git a/packages/coding-agent/src/session/messages.ts b/packages/coding-agent/src/session/messages.ts index 3cba7ca1f..bf21f4f5b 100644 --- a/packages/coding-agent/src/session/messages.ts +++ b/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 { 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 { 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 diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index 8756462f3..e4841533e 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/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 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( @@ -1832,6 +1842,7 @@ export class SessionManager { content: string | (TextContent | ImageContent)[], display: boolean, details?: T, + attribution: MessageAttribution = "agent", ): string { const entry: CustomMessageEntry = { 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(), diff --git a/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts b/packages/coding-agent/test/agent-session-before-agent-start-attribution.test.ts new file mode 100644 index 000000000..ef8c4c226 --- /dev/null +++ b/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"); + }); +}); diff --git a/packages/coding-agent/test/session-manager/build-context.test.ts b/packages/coding-agent/test/session-manager/build-context.test.ts index 6b62446aa..994e8bc27 100644 --- a/packages/coding-agent/test/session-manager/build-context.test.ts +++ b/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"), diff --git a/packages/coding-agent/test/session-messages.test.ts b/packages/coding-agent/test/session-messages.test.ts new file mode 100644 index 000000000..3c6c81d53 --- /dev/null +++ b/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: "Read file", + 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"); + }); +});