From 3a0f80d9a75e7bed7b22745578435d94bf5c580b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 23 May 2026 13:54:57 -0400 Subject: [PATCH] feat(new): enforce chat membership integrity --- .../src/main/services/chat-service.ts | 149 +++++++++++++++--- .../src/main/services/organization-service.ts | 13 ++ .../src/renderer/src/pages/ChatPage.tsx | 14 ++ AgentCompany-New/src/shared/commands.ts | 2 + .../tests/unit/domain-services.test.ts | 124 +++++++++++++++ 5 files changed, 281 insertions(+), 21 deletions(-) diff --git a/AgentCompany-New/src/main/services/chat-service.ts b/AgentCompany-New/src/main/services/chat-service.ts index cdd49e78..160a7301 100644 --- a/AgentCompany-New/src/main/services/chat-service.ts +++ b/AgentCompany-New/src/main/services/chat-service.ts @@ -6,6 +6,23 @@ import { AuditService } from "./audit-service"; const SPACE_TYPES = new Set(["company", "department", "project", "direct"]); const SPACE_VISIBILITIES = new Set(["workspace", "restricted"]); +interface AttachmentCandidate { + fileId: string; + checksum: string; +} + +interface SpaceRow { + id: string; + org_id: string; + visibility: string; + space_type: string; +} + +interface PrincipalRef { + type: "employee" | "agent"; + id: string; +} + export class ChatService { private readonly audit: AuditService; @@ -20,6 +37,9 @@ export class ChatService { if (!topic) throw new Error("Chat space topic is required."); if (!SPACE_TYPES.has(input.spaceType)) throw new Error(`Unsupported chat space type: ${input.spaceType}`); if (!SPACE_VISIBILITIES.has(input.visibility)) throw new Error(`Unsupported chat space visibility: ${input.visibility}`); + const explicitMembers = this.resolveExplicitMembers(input); + const broadSpace = input.visibility === "workspace" && input.spaceType !== "direct"; + if (!broadSpace && explicitMembers.length === 0) throw new Error("Explicit members are required for this chat space."); const id = createId("spc"); const now = new Date().toISOString(); @@ -28,13 +48,8 @@ export class ChatService { values (?, ?, ?, ?, ?, ?, ?) `).run(id, input.orgId, input.spaceType, name, topic, input.visibility, now); - const employees = this.db.prepare("select id from employees where org_id = ?").all(input.orgId) as Array<{ id: string }>; - for (const employee of employees) { - this.db.prepare(` - insert into chat_memberships(space_id, principal_type, principal_id, role, mute_state) - values (?, ?, ?, ?, ?) - `).run(id, "employee", employee.id, input.visibility === "restricted" ? "approver" : "member", "unmuted"); - } + const members = broadSpace ? this.resolveWorkspaceMembers(input.orgId) : explicitMembers; + for (const member of members) this.addMembership(id, member, input.visibility === "restricted" ? "approver" : "member"); this.audit.record({ orgId: input.orgId, @@ -51,6 +66,14 @@ export class ChatService { sendMessage(input: SendMessageInput): string { const body = input.bodyMd.trim(); if (!body) throw new Error("Message body is required."); + const space = this.requireSpace(input.orgId, input.spaceId); + const sender = this.requireSender(input.orgId, input.senderType, input.senderId); + if (sender) this.requireMembership(space.id, sender); + if (input.threadId) this.requireThread(input.orgId, input.spaceId, input.threadId); + const slash = parseSlashInvocation(body); + if (slash && !ALLOWED_SLASH_COMMANDS.has(slash.command)) throw new Error(`Slash command /${slash.command} is unavailable.`); + const attachments = this.requireAttachments(input.orgId, input.attachmentFileIds ?? []); + const id = createId("msg"); const threadId = input.threadId ?? createId("thread"); const now = new Date().toISOString(); @@ -79,8 +102,8 @@ export class ChatService { this.db.prepare("update chat_threads set updated_at = ? where id = ? and org_id = ?").run(now, input.threadId, input.orgId); } this.recordMentions(id, body); - this.recordSlashInvocation(id, body, now); - this.recordAttachments(input.orgId, id, input.attachmentFileIds ?? []); + if (slash) this.recordSlashInvocation(id, slash, now); + this.recordAttachments(id, attachments); if (input.senderId) { this.db.prepare(` insert into chat_presence(principal_type, principal_id, status, status_text, last_seen_at) @@ -108,6 +131,7 @@ export class ChatService { const employee = this.db.prepare("select id from employees where id = ? and org_id = ?") .get(input.employeeId, input.orgId) as { id: string } | undefined; if (!employee) throw new Error("Employee was not found."); + this.requireMembership(input.spaceId, { type: "employee", id: input.employeeId }); const latest = this.db.prepare(` select id from chat_messages where space_id = ? and org_id = ? @@ -138,28 +162,105 @@ export class ChatService { } } - private recordSlashInvocation(messageId: string, body: string, now: string): void { - const match = body.match(/^\/([a-z][a-z0-9_-]*)(?:\s+([\s\S]*))?$/i); - if (!match) return; - const command = match[1].toLowerCase(); - if (!ALLOWED_SLASH_COMMANDS.has(command)) throw new Error(`Slash command /${command} is unavailable.`); + private recordSlashInvocation(messageId: string, slash: SlashInvocation, now: string): void { this.db.prepare(` insert into chat_slash_invocations(id, message_id, command, args_json, result_ref, created_at) values (?, ?, ?, ?, ?, ?) - `).run(createId("slash"), messageId, command, JSON.stringify({ raw: match[2] ?? "" }), null, now); + `).run(createId("slash"), messageId, slash.command, JSON.stringify({ raw: slash.args }), null, now); } - private recordAttachments(orgId: string, messageId: string, fileIds: string[]): void { - for (const fileId of fileIds) { - const file = this.db.prepare("select sha256 from file_objects where id = ? and org_id = ?") - .get(fileId, orgId) as { sha256: string } | undefined; - if (!file) throw new Error("Attachment file was not found."); + private recordAttachments(messageId: string, attachments: AttachmentCandidate[]): void { + for (const file of attachments) { this.db.prepare(` insert into chat_attachments(message_id, file_id, transfer_type, checksum, preview_key) values (?, ?, ?, ?, ?) - `).run(messageId, fileId, "handoff", file.sha256, null); + `).run(messageId, file.fileId, "handoff", file.checksum, null); } } + + private requireSpace(orgId: string, spaceId: string): SpaceRow { + const space = this.db.prepare("select id, org_id, visibility, space_type from chat_spaces where id = ? and org_id = ?") + .get(spaceId, orgId) as SpaceRow | undefined; + if (!space) throw new Error("Chat space was not found."); + return space; + } + + private requireSender(orgId: string, senderType: SendMessageInput["senderType"], senderId: string | null): PrincipalRef | null { + if (senderType === "system") return null; + if (senderType === "employee") { + if (!senderId) throw new Error("Employee sender is required."); + this.requireOrgEntity("employees", senderId, orgId, "Employee sender"); + return { type: "employee", id: senderId }; + } + if (senderType === "agent") { + if (!senderId) throw new Error("Agent sender is required."); + this.requireOrgEntity("agents", senderId, orgId, "Agent sender"); + return { type: "agent", id: senderId }; + } + if (!senderId) throw new Error("Board sender is required."); + this.requireOrgEntity("employees", senderId, orgId, "Board sender"); + return { type: "employee", id: senderId }; + } + + private requireThread(orgId: string, spaceId: string, threadId: string): void { + const thread = this.db.prepare("select id from chat_threads where id = ? and org_id = ? and space_id = ?") + .get(threadId, orgId, spaceId) as { id: string } | undefined; + if (!thread) throw new Error("Chat thread was not found in this space."); + } + + private requireAttachments(orgId: string, fileIds: string[]): AttachmentCandidate[] { + return fileIds.map((fileId) => { + const file = this.db.prepare("select sha256 from file_objects where id = ? and org_id = ?") + .get(fileId, orgId) as { sha256: string } | undefined; + if (!file) throw new Error("Attachment file was not found."); + return { fileId, checksum: file.sha256 }; + }); + } + + private requireMembership(spaceId: string, principal: PrincipalRef): void { + const row = this.db.prepare("select space_id from chat_memberships where space_id = ? and principal_type = ? and principal_id = ?") + .get(spaceId, principal.type, principal.id) as { space_id: string } | undefined; + if (!row) throw new Error("Chat membership is required."); + } + + private resolveWorkspaceMembers(orgId: string): PrincipalRef[] { + const employees = this.db.prepare("select id from employees where org_id = ?").all(orgId) as Array<{ id: string }>; + const agents = this.db.prepare("select id from agents where org_id = ?").all(orgId) as Array<{ id: string }>; + return [ + ...employees.map((employee) => ({ type: "employee" as const, id: employee.id })), + ...agents.map((agent) => ({ type: "agent" as const, id: agent.id })), + ]; + } + + private resolveExplicitMembers(input: CreateChatSpaceInput): PrincipalRef[] { + const members: PrincipalRef[] = []; + for (const id of new Set(input.memberEmployeeIds ?? [])) { + this.requireOrgEntity("employees", id, input.orgId, "Chat member"); + members.push({ type: "employee", id }); + } + for (const id of new Set(input.memberAgentIds ?? [])) { + this.requireOrgEntity("agents", id, input.orgId, "Chat member"); + members.push({ type: "agent", id }); + } + return members; + } + + private requireOrgEntity(table: "employees" | "agents", id: string, orgId: string, label: string): void { + const row = this.db.prepare(`select id from ${table} where id = ? and org_id = ?`).get(id, orgId) as { id: string } | undefined; + if (!row) throw new Error(`${label} was not found.`); + } + + private addMembership(spaceId: string, principal: PrincipalRef, role: string): void { + this.db.prepare(` + insert into chat_memberships(space_id, principal_type, principal_id, role, mute_state) + values (?, ?, ?, ?, ?) + `).run(spaceId, principal.type, principal.id, role, "unmuted"); + } +} + +interface SlashInvocation { + command: string; + args: string; } const ALLOWED_SLASH_COMMANDS = new Set([ @@ -173,6 +274,12 @@ const ALLOWED_SLASH_COMMANDS = new Set([ "escalate", ]); +function parseSlashInvocation(body: string): SlashInvocation | null { + const match = body.match(/^\/([a-z][a-z0-9_-]*)(?:\s+([\s\S]*))?$/i); + if (!match) return null; + return { command: match[1].toLowerCase(), args: match[2] ?? "" }; +} + function createThreadTitle(body: string): string { const firstLine = body.split(/\r?\n/, 1)[0].trim(); return firstLine.length > 80 ? `${firstLine.slice(0, 77)}...` : firstLine; diff --git a/AgentCompany-New/src/main/services/organization-service.ts b/AgentCompany-New/src/main/services/organization-service.ts index 7d88fd3b..68c49eaf 100644 --- a/AgentCompany-New/src/main/services/organization-service.ts +++ b/AgentCompany-New/src/main/services/organization-service.ts @@ -163,6 +163,7 @@ export class OrganizationService { action: "created", payload: { role: input.role, department: input.department }, }); + this.addAgentToWorkspaceSpaces(input.orgId, id); return id; } @@ -244,6 +245,18 @@ export class OrganizationService { if (goalId && project.goal_id !== goalId) throw new Error("Project is linked to a different goal."); } + private addAgentToWorkspaceSpaces(orgId: string, agentId: string): void { + const spaces = this.db.prepare("select id from chat_spaces where org_id = ? and visibility = ? and space_type <> ?") + .all(orgId, "workspace", "direct") as Array<{ id: string }>; + for (const space of spaces) { + this.db.prepare(` + insert into chat_memberships(space_id, principal_type, principal_id, role, mute_state) + values (?, ?, ?, ?, ?) + on conflict(space_id, principal_type, principal_id) do nothing + `).run(space.id, "agent", agentId, "member", "unmuted"); + } + } + updateTaskStatus(orgId: string, taskId: string, status: string): void { const result = this.db.prepare("update tasks set status = ?, updated_at = ? where id = ? and org_id = ?") .run(status, new Date().toISOString(), taskId, orgId); diff --git a/AgentCompany-New/src/renderer/src/pages/ChatPage.tsx b/AgentCompany-New/src/renderer/src/pages/ChatPage.tsx index d8b4532d..c2addc33 100644 --- a/AgentCompany-New/src/renderer/src/pages/ChatPage.tsx +++ b/AgentCompany-New/src/renderer/src/pages/ChatPage.tsx @@ -127,6 +127,8 @@ export function ChatPage({ snapshot, actions }: PageProps) { name: value(form, "name"), topic: value(form, "topic"), visibility: value(form, "visibility") as Visibility, + memberEmployeeIds: [employee.id, ...form.getAll("memberEmployeeIds").map(String)].filter(Boolean), + memberAgentIds: form.getAll("memberAgentIds").map(String).filter(Boolean), })).then((result) => { if (result) setCreateOpen(false); }); }; @@ -207,6 +209,18 @@ export function ChatPage({ snapshot, actions }: PageProps) { + + + + + + ); diff --git a/AgentCompany-New/src/shared/commands.ts b/AgentCompany-New/src/shared/commands.ts index 8ad861ba..61073339 100644 --- a/AgentCompany-New/src/shared/commands.ts +++ b/AgentCompany-New/src/shared/commands.ts @@ -95,6 +95,8 @@ export interface CreateChatSpaceInput { name: string; topic: string; visibility: "workspace" | "restricted"; + memberEmployeeIds?: string[]; + memberAgentIds?: string[]; } export interface MarkChatSpaceReadInput { diff --git a/AgentCompany-New/tests/unit/domain-services.test.ts b/AgentCompany-New/tests/unit/domain-services.test.ts index aceb0bc7..f38b4a41 100644 --- a/AgentCompany-New/tests/unit/domain-services.test.ts +++ b/AgentCompany-New/tests/unit/domain-services.test.ts @@ -237,6 +237,130 @@ describe("domain services", () => { connection.close(); }); + test("rejects chat writes outside membership and prevents partial invalid messages", () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + const chat = new ChatService(connection.db); + let orgId = ""; + let employeeId = ""; + let agentId = ""; + let restrictedSpaceId = ""; + let workspaceSpaceId = ""; + + connection.transaction(() => { + orgId = organization.createOrganization({ + organizationName: "Acme Operations", + employeeName: "Dana Lee", + mission: "Coordinate agents with auditable approvals.", + goal: "Ship a verified control plane.", + }); + employeeId = (connection.db.prepare("select id from employees where org_id = ?").get(orgId) as { id: string }).id; + agentId = organization.createAgent({ + orgId, + name: "Ada", + role: "Release reviewer", + department: "Engineering", + modelProviderId: null, + modelId: null, + budgetCents: 10000, + managerAgentId: null, + }); + restrictedSpaceId = chat.createSpace({ + orgId, + spaceType: "project", + name: "board-review", + topic: "Restricted release review", + visibility: "restricted", + memberEmployeeIds: [employeeId], + }); + workspaceSpaceId = chat.createSpace({ + orgId, + spaceType: "project", + name: "release", + topic: "Release coordination", + visibility: "workspace", + }); + }); + + expect(() => connection.transaction(() => chat.sendMessage({ + orgId, + spaceId: restrictedSpaceId, + threadId: null, + senderType: "agent", + senderId: agentId, + bodyMd: "I should not be here.", + attachmentFileIds: [], + }))).toThrow("Chat membership is required."); + expect(() => connection.transaction(() => chat.sendMessage({ + orgId, + spaceId: workspaceSpaceId, + threadId: null, + senderType: "employee", + senderId: employeeId, + bodyMd: "/unknown command", + attachmentFileIds: [], + }))).toThrow("Slash command /unknown is unavailable."); + expect(() => connection.transaction(() => chat.sendMessage({ + orgId, + spaceId: workspaceSpaceId, + threadId: null, + senderType: "employee", + senderId: employeeId, + bodyMd: "Missing file", + attachmentFileIds: ["file_missing"], + }))).toThrow("Attachment file was not found."); + + const count = connection.db.prepare("select count(*) as count from chat_messages").get() as { count: number }; + expect(count.count).toBe(0); + connection.close(); + }); + + test("rejects replies to threads from another chat space", () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + const chat = new ChatService(connection.db); + let orgId = ""; + let employeeId = ""; + let firstSpaceId = ""; + let secondSpaceId = ""; + let threadId = ""; + + connection.transaction(() => { + orgId = organization.createOrganization({ + organizationName: "Acme Operations", + employeeName: "Dana Lee", + mission: "Coordinate agents with auditable approvals.", + goal: "Ship a verified control plane.", + }); + employeeId = (connection.db.prepare("select id from employees where org_id = ?").get(orgId) as { id: string }).id; + firstSpaceId = chat.createSpace({ orgId, spaceType: "project", name: "release-a", topic: "A", visibility: "workspace" }); + secondSpaceId = chat.createSpace({ orgId, spaceType: "project", name: "release-b", topic: "B", visibility: "workspace" }); + chat.sendMessage({ + orgId, + spaceId: firstSpaceId, + threadId: null, + senderType: "employee", + senderId: employeeId, + bodyMd: "Root", + attachmentFileIds: [], + }); + threadId = (connection.db.prepare("select id from chat_threads where space_id = ?").get(firstSpaceId) as { id: string }).id; + }); + + expect(() => connection.transaction(() => chat.sendMessage({ + orgId, + spaceId: secondSpaceId, + threadId, + senderType: "employee", + senderId: employeeId, + bodyMd: "Wrong thread", + attachmentFileIds: [], + }))).toThrow("Chat thread was not found in this space."); + const count = connection.db.prepare("select count(*) as count from chat_messages").get() as { count: number }; + expect(count.count).toBe(1); + connection.close(); + }); + test("creates a read-only migration report from a legacy database", () => { const connection = createConnection(); const organization = new OrganizationService(connection.db);