From fd7f9b9ff7f40ffd2a7b00141e164d35e7e1a850 Mon Sep 17 00:00:00 2001 From: test Date: Sun, 24 May 2026 17:25:12 +0800 Subject: [PATCH 1/2] Add trusted multi-user chat prefixing --- relay/src/relay-object.ts | 3 +- relay/src/sharing.test.ts | 11 ++ relay/src/sharing.ts | 9 ++ src/collaboration.test.ts | 86 ++++++++++++++++ src/collaboration.ts | 100 +++++++++++++++++++ src/metadata-server.test.ts | 73 ++++++++++++++ src/metadata-server.ts | 5 +- src/multiplexer/agent-io-methods.ts | 4 +- src/multiplexer/dashboard-model.ts | 9 +- src/multiplexer/session-runtime-core.test.ts | 51 ++++++++++ src/multiplexer/session-runtime-core.ts | 16 +-- src/session-message-history.ts | 13 +++ 12 files changed, 368 insertions(+), 12 deletions(-) create mode 100644 src/collaboration.test.ts create mode 100644 src/collaboration.ts diff --git a/relay/src/relay-object.ts b/relay/src/relay-object.ts index 2b012c74..887f615b 100644 --- a/relay/src/relay-object.ts +++ b/relay/src/relay-object.ts @@ -25,6 +25,7 @@ import { saveSharingState, sharedRelayRequestAccess, summarizeShare, + stripTrustedAimuxHeaders, } from "./sharing.js"; const HEARTBEAT_INTERVAL_MS = 30_000; @@ -460,7 +461,7 @@ export class RelayObject extends DurableObject { ok: true, requestPatch: { headers: { - ...(request.headers ?? {}), + ...stripTrustedAimuxHeaders(request.headers), "X-Aimux-Share-Id": share.id, "X-Aimux-Share-Mode": getShareChatMode(share), "X-Aimux-Actor-User-Id": participant.userId, diff --git a/relay/src/sharing.test.ts b/relay/src/sharing.test.ts index 61909a3a..4ade8453 100644 --- a/relay/src/sharing.test.ts +++ b/relay/src/sharing.test.ts @@ -8,6 +8,7 @@ import { isSharedRelayRequestAllowed, removeShareParticipant, sharedRelayRequestAccess, + stripTrustedAimuxHeaders, summarizeShare, } from "./sharing"; @@ -190,6 +191,16 @@ describe("sharing state", () => { expect(actorDisplayPrefix({ userId: "u", displayName: "", role: "guest" })).toBe("[User]:"); }); + it("strips client-provided trusted relay headers before injection", () => { + expect( + stripTrustedAimuxHeaders({ + "content-type": "application/json", + "x-aimux-actor-name": "Mallory", + "X-Aimux-Share-Mode": "multi", + }), + ).toEqual({ "content-type": "application/json" }); + }); + it("redacts invite token hashes from public summaries", async () => { const created = await createShareInvite(emptySharingState(), { owner, diff --git a/relay/src/sharing.ts b/relay/src/sharing.ts index 65aec53c..569a628f 100644 --- a/relay/src/sharing.ts +++ b/relay/src/sharing.ts @@ -263,6 +263,15 @@ export function actorDisplayPrefix(actor: ShareActor): string { return `[${sanitizeDisplayName(actor.displayName)}]:`; } +export function stripTrustedAimuxHeaders(headers: Record | undefined): Record { + const safeHeaders: Record = {}; + for (const [key, value] of Object.entries(headers ?? {})) { + if (key.toLowerCase().startsWith("x-aimux-")) continue; + safeHeaders[key] = value; + } + return safeHeaders; +} + function createShare(input: { owner: ShareActor; projectRoot: string; diff --git a/src/collaboration.test.ts b/src/collaboration.test.ts new file mode 100644 index 00000000..a847d4b5 --- /dev/null +++ b/src/collaboration.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, it } from "vitest"; + +import { applyAgentCollaborationPrefix, collaborationContextFromHeaders } from "./collaboration.js"; + +describe("collaboration context", () => { + it("parses relay-injected actor headers", () => { + expect( + collaborationContextFromHeaders({ + "x-aimux-share-id": "share_123", + "x-aimux-share-mode": "multi", + "x-aimux-actor-user-id": "user_123", + "x-aimux-actor-name": "Sam Steady", + "x-aimux-actor-email": "sam@example.com", + "x-aimux-actor-role": "owner", + }), + ).toEqual({ + shareId: "share_123", + mode: "multi", + actor: { + userId: "user_123", + displayName: "Sam Steady", + email: "sam@example.com", + role: "owner", + }, + }); + }); + + it("returns undefined when no collaboration headers are present", () => { + expect(collaborationContextFromHeaders({ "content-type": "application/json" })).toBeUndefined(); + }); + + it("prefixes plain text input only in multi-user mode", () => { + const collaboration = { + shareId: "share_123", + mode: "multi" as const, + actor: { userId: "user_123", displayName: "Sam Steady" }, + }; + + expect(applyAgentCollaborationPrefix({ data: "ship it" }, collaboration)).toEqual({ + data: "[Sam Steady]: ship it", + parts: undefined, + }); + expect(applyAgentCollaborationPrefix({ data: "ship it" }, { ...collaboration, mode: "single" })).toEqual({ + data: "ship it", + }); + }); + + it("prefixes the first text part without mutating image parts", () => { + const input = { + parts: [ + { type: "image" as const, attachmentId: "att_1", alt: "screenshot" }, + { type: "text" as const, text: "look here" }, + ], + }; + + expect( + applyAgentCollaborationPrefix(input, { + mode: "multi", + actor: { userId: "user_123", displayName: "Casey" }, + }), + ).toEqual({ + parts: [ + { type: "image", attachmentId: "att_1", alt: "screenshot" }, + { type: "text", text: "[Casey]: look here" }, + ], + }); + }); + + it("adds a speaker-only text part before image-only input", () => { + expect( + applyAgentCollaborationPrefix( + { parts: [{ type: "image", attachmentId: "att_1" }] }, + { + mode: "multi", + actor: { userId: "user_123", displayName: "Casey" }, + }, + ), + ).toEqual({ + data: undefined, + parts: [ + { type: "text", text: "[Casey]:" }, + { type: "image", attachmentId: "att_1" }, + ], + }); + }); +}); diff --git a/src/collaboration.ts b/src/collaboration.ts new file mode 100644 index 00000000..b87b564a --- /dev/null +++ b/src/collaboration.ts @@ -0,0 +1,100 @@ +import type { IncomingHttpHeaders } from "node:http"; + +import type { AgentInputPart } from "./agent-message-parts.js"; + +export type CollaborationChatMode = "single" | "multi"; +export type CollaborationRole = "owner" | "guest"; + +export interface CollaborationActor { + userId: string; + displayName: string; + email?: string; + role?: CollaborationRole; +} + +export interface AgentCollaborationContext { + shareId?: string; + mode?: CollaborationChatMode; + actor?: CollaborationActor; +} + +type HeaderMap = IncomingHttpHeaders | Record; + +export function collaborationContextFromHeaders(headers: HeaderMap): AgentCollaborationContext | undefined { + const shareId = headerValue(headers, "x-aimux-share-id"); + const mode = parseMode(headerValue(headers, "x-aimux-share-mode")); + const userId = headerValue(headers, "x-aimux-actor-user-id"); + const displayName = headerValue(headers, "x-aimux-actor-name"); + const email = headerValue(headers, "x-aimux-actor-email"); + const role = parseRole(headerValue(headers, "x-aimux-actor-role")); + + if (!shareId && !mode && !userId && !displayName && !email && !role) { + return undefined; + } + + return { + shareId, + mode, + actor: userId + ? { + userId, + displayName: displayName || email || userId, + email, + role, + } + : undefined, + }; +} + +export function applyAgentCollaborationPrefix( + input: { data?: string; parts?: AgentInputPart[] }, + collaboration?: AgentCollaborationContext, +): { data?: string; parts?: AgentInputPart[] } { + const actor = collaboration?.actor; + if (collaboration?.mode !== "multi" || !actor) { + return input; + } + + const prefix = `[${formatSpeakerName(actor.displayName)}]:`; + if (Array.isArray(input.parts) && input.parts.length > 0) { + let applied = false; + const parts = input.parts.map((part) => { + if (part.type !== "text" || applied) return part; + const text = String(part.text ?? ""); + if (!text.trim()) return part; + applied = true; + return { ...part, text: `${prefix} ${text}` }; + }); + return { + data: input.data, + parts: applied ? parts : [{ type: "text", text: prefix }, ...input.parts], + }; + } + + const data = String(input.data ?? ""); + return { + data: data.trim() ? `${prefix} ${data}` : prefix, + parts: input.parts, + }; +} + +function headerValue(headers: HeaderMap, name: string): string | undefined { + const raw = headers[name] ?? headers[name.toLowerCase()]; + const value = Array.isArray(raw) ? raw[0] : raw; + const trimmed = value?.trim(); + return trimmed || undefined; +} + +function parseMode(value?: string): CollaborationChatMode | undefined { + const mode = value?.toLowerCase(); + return mode === "single" || mode === "multi" ? mode : undefined; +} + +function parseRole(value?: string): CollaborationRole | undefined { + const role = value?.toLowerCase(); + return role === "owner" || role === "guest" ? role : undefined; +} + +function formatSpeakerName(name: string): string { + return name.replace(/\s+/g, " ").trim().slice(0, 80) || "User"; +} diff --git a/src/metadata-server.test.ts b/src/metadata-server.test.ts index 2e549c1e..85627a0b 100644 --- a/src/metadata-server.test.ts +++ b/src/metadata-server.test.ts @@ -1582,6 +1582,79 @@ describe("MetadataServer threads API", () => { expect(writes).toEqual([{ sessionId: "codex-1", data: "hello from http", submit: true }]); }); + it("passes trusted collaboration headers with agent input over HTTP", async () => { + server?.stop(); + const writes: Array<{ + sessionId: string; + data?: string; + collaboration?: { + shareId?: string; + mode?: string; + actor?: { userId: string; displayName: string; email?: string; role?: string }; + }; + }> = []; + server = new MetadataServer({ + lifecycle: { + writeAgentInput: ({ sessionId, data, collaboration }) => { + writes.push({ sessionId, data, collaboration }); + return { + sessionId, + accepted: true, + operation: { + id: "inputop-collab", + sessionId, + submit: true, + state: "submitted" as const, + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + }, + }; + }, + }, + }); + await server.start(); + + const endpoint = server?.getAddress(); + expect(endpoint).toBeTruthy(); + const base = `http://${endpoint!.host}:${endpoint!.port}`; + + const inputRes = await fetch(`${base}/agents/input`, { + method: "POST", + headers: { + "content-type": "application/json", + "x-aimux-share-id": "share_123", + "x-aimux-share-mode": "multi", + "x-aimux-actor-user-id": "user_123", + "x-aimux-actor-name": "Sam Steady", + "x-aimux-actor-email": "sam@example.com", + "x-aimux-actor-role": "owner", + }, + body: JSON.stringify({ + sessionId: "claude-1", + data: "hello from http", + submit: true, + }), + }); + + expect(inputRes.ok).toBe(true); + expect(writes).toEqual([ + { + sessionId: "claude-1", + data: "hello from http", + collaboration: { + shareId: "share_123", + mode: "multi", + actor: { + userId: "user_123", + displayName: "Sam Steady", + email: "sam@example.com", + role: "owner", + }, + }, + }, + ]); + }); + it("passes structured input parts with agent input over HTTP", async () => { server?.stop(); const writes: Array<{ diff --git a/src/metadata-server.ts b/src/metadata-server.ts index 9d675988..35a26912 100644 --- a/src/metadata-server.ts +++ b/src/metadata-server.ts @@ -3,6 +3,7 @@ import { createHash, randomUUID } from "node:crypto"; import { mkdirSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs"; import { dirname, join } from "node:path"; import { getDashboardClientUiStatePath, getPlansDir, getProjectId, getProjectStateDir } from "./paths.js"; +import { collaborationContextFromHeaders, type AgentCollaborationContext } from "./collaboration.js"; import { type MetadataTone, updateSessionMetadata, @@ -282,6 +283,7 @@ interface MetadataServerOptions { parts?: AgentInputPart[]; clientMessageId?: string; submit?: boolean; + collaboration?: AgentCollaborationContext; }) => | Promise<{ sessionId: string; @@ -2484,7 +2486,8 @@ export class MetadataServer { send(res, 501, { ok: false, error: "agent input not supported by this service" }); return; } - const result = await this.options.lifecycle.writeAgentInput(body); + const collaboration = collaborationContextFromHeaders(req.headers); + const result = await this.options.lifecycle.writeAgentInput({ ...body, collaboration }); if (this.options.lifecycle.readAgentHistory) { try { const history = await this.options.lifecycle.readAgentHistory({ sessionId: body.sessionId, lastN: 20 }); diff --git a/src/multiplexer/agent-io-methods.ts b/src/multiplexer/agent-io-methods.ts index b516be7c..9811c831 100644 --- a/src/multiplexer/agent-io-methods.ts +++ b/src/multiplexer/agent-io-methods.ts @@ -2,6 +2,7 @@ import { sendDirectMessage, sendThreadMessage } from "../orchestration.js"; import { sendHandoff } from "../orchestration-actions.js"; import { resolveOrchestrationRecipients } from "../orchestration-routing.js"; import type { DashboardSession } from "../dashboard/index.js"; +import type { AgentCollaborationContext } from "../collaboration.js"; import { markMessageDelivered, type MessageKind } from "../threads.js"; import { stopProjectServices as stopProjectServicesImpl } from "./dashboard-model.js"; import { @@ -257,8 +258,9 @@ export const agentIoMethods = { parts?: any[], clientMessageId?: string, submit = false, + collaboration?: AgentCollaborationContext, ): Promise<{ sessionId: string }> { - return writeAgentInputImpl(this, sessionId, data, parts, clientMessageId, submit); + return writeAgentInputImpl(this, sessionId, data, parts, clientMessageId, submit, collaboration); }, async readAgentHistory(this: any, sessionId: string, lastN?: number): Promise { diff --git a/src/multiplexer/dashboard-model.ts b/src/multiplexer/dashboard-model.ts index 54db0da9..aebbd72b 100644 --- a/src/multiplexer/dashboard-model.ts +++ b/src/multiplexer/dashboard-model.ts @@ -1188,7 +1188,14 @@ export async function startProjectServices(host: DashboardModelHost): Promise host.recordSessionBackendSessionId(input.sessionId, input.backendSessionId), writeAgentInput: (input: any) => - host.writeAgentInput(input.sessionId, input.data, input.parts, input.clientMessageId, input.submit), + host.writeAgentInput( + input.sessionId, + input.data, + input.parts, + input.clientMessageId, + input.submit, + input.collaboration, + ), readAgentOutput: (input: any) => host.readAgentOutput(input.sessionId, input.startLine), readAgentHistory: (input: any) => host.readAgentHistory(input.sessionId, input.lastN), }, diff --git a/src/multiplexer/session-runtime-core.test.ts b/src/multiplexer/session-runtime-core.test.ts index 3e5563ba..d57bfc3b 100644 --- a/src/multiplexer/session-runtime-core.test.ts +++ b/src/multiplexer/session-runtime-core.test.ts @@ -6,6 +6,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { initPaths } from "../paths.js"; import { recordSessionBackendSessionIdMetadata } from "../metadata-store.js"; import { readSessionInputOperation } from "../session-input-operations.js"; +import { readSessionMessages } from "../session-message-history.js"; import { TmuxSessionTransport } from "../tmux/session-transport.js"; import { buildTmuxWindowMetadata, @@ -175,6 +176,56 @@ describe("session runtime prompt submission", () => { } }); + it("prefixes multi-user input for the agent while storing original message metadata", async () => { + const repoRoot = mkdtempSync(join(tmpdir(), "aimux-session-runtime-")); + try { + await initPaths(repoRoot); + const writes: string[] = []; + const host: any = { + sessions: [ + { + id: "claude-1", + exited: false, + write: (data: string) => writes.push(data), + }, + ], + sessionToolKeys: new Map([["claude-1", "claude"]]), + }; + + const result = await writeAgentInput(host, "claude-1", "Can you check this?", undefined, "client-1", true, { + shareId: "share_123", + mode: "multi", + actor: { + userId: "user_123", + displayName: "Sam Steady", + email: "sam@example.com", + role: "owner", + }, + }); + + expect(result.accepted).toBe(true); + expect(writes).toEqual(["[Sam Steady]: Can you check this?\r"]); + expect(readSessionMessages("claude-1")).toMatchObject([ + { + clientMessageId: "client-1", + sessionId: "claude-1", + role: "user", + parts: [{ type: "text", text: "Can you check this?" }], + actor: { + userId: "user_123", + displayName: "Sam Steady", + email: "sam@example.com", + role: "owner", + }, + shareId: "share_123", + chatMode: "multi", + }, + ]); + } finally { + rmSync(repoRoot, { recursive: true, force: true }); + } + }); + it("fails submitted tmux input when the target disappears before prompt submission", async () => { const repoRoot = mkdtempSync(join(tmpdir(), "aimux-session-runtime-")); vi.useFakeTimers(); diff --git a/src/multiplexer/session-runtime-core.ts b/src/multiplexer/session-runtime-core.ts index d18af36f..8f6cd1f4 100644 --- a/src/multiplexer/session-runtime-core.ts +++ b/src/multiplexer/session-runtime-core.ts @@ -11,6 +11,7 @@ import { loadMetadataState } from "../metadata-store.js"; import { parseAgentOutput } from "../agent-output-parser.js"; import { serializeAgentInput } from "../agent-message-parts.js"; import { resolveAttachmentPath } from "../attachment-store.js"; +import { applyAgentCollaborationPrefix, type AgentCollaborationContext } from "../collaboration.js"; import { appendSessionMessage, readSessionMessages } from "../session-message-history.js"; import { createSessionInputOperation, @@ -254,6 +255,7 @@ export async function writeAgentInput( parts?: any[], clientMessageId?: string, submit = false, + collaboration?: AgentCollaborationContext, ): Promise<{ sessionId: string; accepted: boolean; @@ -262,13 +264,11 @@ export async function writeAgentInput( error?: string; }> { const session = resolveRunningSession(host, sessionId); - const serializedData = serializeAgentInput( - { data, parts }, - { - tool: host.sessionToolKeys.get(sessionId), - resolveAttachmentPath, - }, - ); + const agentInput = applyAgentCollaborationPrefix({ data, parts }, collaboration); + const serializedData = serializeAgentInput(agentInput, { + tool: host.sessionToolKeys.get(sessionId), + resolveAttachmentPath, + }); const normalizedData = normalizeAgentInput(host, serializedData, submit, sessionId); if (!normalizedData && !submit) { throw new Error("input data is required"); @@ -276,7 +276,7 @@ export async function writeAgentInput( let operation = createSessionInputOperation({ sessionId, clientMessageId, submit }); try { - const message = appendSessionMessage(sessionId, { data, parts, clientMessageId }); + const message = appendSessionMessage(sessionId, { data, parts, clientMessageId, collaboration }); if (message?.id) { operation = saveSessionInputOperation({ ...operation, diff --git a/src/session-message-history.ts b/src/session-message-history.ts index 6b09cf0b..ed1ac8fa 100644 --- a/src/session-message-history.ts +++ b/src/session-message-history.ts @@ -4,6 +4,7 @@ import { randomUUID } from "node:crypto"; import { getSessionMessagesDir } from "./paths.js"; import type { AgentInputPart } from "./agent-message-parts.js"; import { getAttachment } from "./attachment-store.js"; +import type { AgentCollaborationContext, CollaborationChatMode, CollaborationRole } from "./collaboration.js"; export interface SessionMessagePart { type: "text" | "image"; @@ -22,6 +23,14 @@ export interface SessionMessageRecord { role: "user"; ts: string; parts: SessionMessagePart[]; + actor?: { + userId: string; + displayName: string; + email?: string; + role?: CollaborationRole; + }; + shareId?: string; + chatMode?: CollaborationChatMode; } function historyPath(sessionId: string): string { @@ -34,6 +43,7 @@ export function appendSessionMessage( data?: string; parts?: AgentInputPart[]; clientMessageId?: string; + collaboration?: AgentCollaborationContext; }, ): SessionMessageRecord | null { const parts = normalizeMessageParts(input); @@ -49,6 +59,9 @@ export function appendSessionMessage( role: "user", ts: new Date().toISOString(), parts, + actor: input.collaboration?.actor, + shareId: input.collaboration?.shareId, + chatMode: input.collaboration?.mode, }; appendFileSync(historyPath(sessionId), `${JSON.stringify(record)}\n`); return record; From 841157502a7f89307f41761cc09662122314e364 Mon Sep 17 00:00:00 2001 From: test Date: Sun, 24 May 2026 17:32:20 +0800 Subject: [PATCH 2/2] Fix collaboration header lookup --- src/collaboration.test.ts | 12 ++++++------ src/collaboration.ts | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/collaboration.test.ts b/src/collaboration.test.ts index a847d4b5..fbeeb80e 100644 --- a/src/collaboration.test.ts +++ b/src/collaboration.test.ts @@ -6,12 +6,12 @@ describe("collaboration context", () => { it("parses relay-injected actor headers", () => { expect( collaborationContextFromHeaders({ - "x-aimux-share-id": "share_123", - "x-aimux-share-mode": "multi", - "x-aimux-actor-user-id": "user_123", - "x-aimux-actor-name": "Sam Steady", - "x-aimux-actor-email": "sam@example.com", - "x-aimux-actor-role": "owner", + "X-Aimux-Share-Id": "share_123", + "X-Aimux-Share-Mode": "multi", + "X-Aimux-Actor-User-Id": "user_123", + "X-Aimux-Actor-Name": "Sam Steady", + "X-Aimux-Actor-Email": "sam@example.com", + "X-Aimux-Actor-Role": "owner", }), ).toEqual({ shareId: "share_123", diff --git a/src/collaboration.ts b/src/collaboration.ts index b87b564a..75f31f9b 100644 --- a/src/collaboration.ts +++ b/src/collaboration.ts @@ -79,7 +79,8 @@ export function applyAgentCollaborationPrefix( } function headerValue(headers: HeaderMap, name: string): string | undefined { - const raw = headers[name] ?? headers[name.toLowerCase()]; + const target = name.toLowerCase(); + const raw = Object.entries(headers).find(([key]) => key.toLowerCase() === target)?.[1]; const value = Array.isArray(raw) ? raw[0] : raw; const trimmed = value?.trim(); return trimmed || undefined;