diff --git a/src/main/services/agent-attachment-tools.ts b/src/main/services/agent-attachment-tools.ts new file mode 100644 index 00000000..fcee5ed5 --- /dev/null +++ b/src/main/services/agent-attachment-tools.ts @@ -0,0 +1,54 @@ +import type { DatabaseSync } from "node:sqlite"; +import { readFileSync } from "node:fs"; +import { FILE_SCAN_STATE } from "../../shared/domain"; +import { requireString } from "./agent-tool-args"; +import type { AgentToolContext, AgentToolDeps } from "./agent-tools"; +import type { JsonRecord } from "./provider-http"; +import type { ToolDefinition } from "./provider-tool-turn"; + +const ATTACHMENT_CHAR_LIMIT = 24_000; + +export const ATTACHMENT_TOOLS: ToolDefinition[] = [ + { + name: "read_attachment", + description: "Read a file attached to a message or handed off to you, by the file_id shown next to the message.", + parameters: { + type: "object", + additionalProperties: false, + properties: { file_id: { type: "string" } }, + required: ["file_id"], + }, + }, +]; + +/** " [attachments: a.md (file_id=…)]" for a message, or "". */ +export function attachmentNote(db: DatabaseSync, messageId: string): string { + const files = db.prepare(` + select f.id, f.name from chat_attachments ca join file_objects f on f.id = ca.file_id + where ca.message_id = ? order by f.name + `).all(messageId) as Array<{ id: string; name: string }>; + if (files.length === 0) return ""; + return ` [attachments: ${files.map((file) => `${file.name} (file_id=${file.id})`).join(", ")}]`; +} + +export function readAttachmentTool(deps: AgentToolDeps, ctx: AgentToolContext, args: JsonRecord): string { + const fileId = requireString(args.file_id, "file_id"); + const file = deps.db.prepare("select name, storage_path, scan_state from file_objects where id = ? and org_id = ?") + .get(fileId, ctx.orgId) as { name: string; storage_path: string; scan_state: string } | undefined; + if (!file) throw new Error(`No file with file_id "${fileId}".`); + if (file.scan_state !== FILE_SCAN_STATE.APPROVED) throw new Error(`${file.name} was blocked by the file scan and cannot be read.`); + const bytes = readFileSync(file.storage_path); + const content = utf8Text(bytes); + if (content === null) return `${file.name} is a binary file (${bytes.length} bytes); it has no text to read.`; + if (content.length <= ATTACHMENT_CHAR_LIMIT) return content; + return `${content.slice(0, ATTACHMENT_CHAR_LIMIT)}\n\n[truncated: showing the first ${ATTACHMENT_CHAR_LIMIT} of ${content.length} characters]`; +} + +function utf8Text(bytes: Buffer): string | null { + if (bytes.includes(0)) return null; + try { + return new TextDecoder("utf-8", { fatal: true }).decode(bytes); + } catch { + return null; + } +} diff --git a/src/main/services/agent-read-messages.ts b/src/main/services/agent-read-messages.ts index 15ae4d2a..fc929d65 100644 --- a/src/main/services/agent-read-messages.ts +++ b/src/main/services/agent-read-messages.ts @@ -1,3 +1,4 @@ +import { attachmentNote } from "./agent-attachment-tools"; import { optionalString } from "./agent-tool-args"; import type { AgentToolContext, AgentToolDeps } from "./agent-tools"; import type { JsonRecord } from "./provider-http"; @@ -75,14 +76,14 @@ function messagePage( ) as Array<{ id: string; row_id: number; body_md: string; created_at: string; thread_id: string; sender_name: string }>; if (rows.length === 0) return `No messages${filter.query ? ` matching "${filter.query}"` : ""}${filter.before ? " before that point" : ""} in ${label}.`; const page = rows.slice(0, filter.limit).reverse(); - const lines = page.map((m) => `[${m.created_at}] ${m.sender_name} (thread=${m.thread_id}): ${m.body_md}`); + const lines = page.map((m) => `[${m.created_at}] ${m.sender_name} (thread=${m.thread_id}): ${m.body_md}${attachmentNote(deps.db, m.id)}`); if (rows.length > filter.limit) lines.unshift(`(Older messages exist: pass before="${page[0].id}" to read them.)`); return lines.join("\n"); } function unansweredMentions(deps: AgentToolDeps, ctx: AgentToolContext, limit: number): string { const rows = deps.db.prepare(` - select m.body_md, + select m.id, m.body_md, coalesce(a.name, e.display_name, m.sender_type) as sender_name, m.thread_id, s.name as channel_name, s.space_type from chat_mentions cm @@ -98,11 +99,11 @@ function unansweredMentions(deps: AgentToolDeps, ctx: AgentToolContext, limit: n ) order by m.created_at desc limit ? - `).all(ctx.agent.id, ctx.orgId, limit) as Array<{ body_md: string; sender_name: string; thread_id: string; channel_name: string; space_type: string }>; + `).all(ctx.agent.id, ctx.orgId, limit) as Array<{ id: string; body_md: string; sender_name: string; thread_id: string; channel_name: string; space_type: string }>; if (rows.length === 0) return "No unanswered mentions."; return rows.reverse().map((m) => { const where = m.space_type === "direct" ? "in your direct conversation" : `in #${m.channel_name}`; - return `[${m.sender_name} ${where}] (thread=${m.thread_id}): ${m.body_md}`; + return `[${m.sender_name} ${where}] (thread=${m.thread_id}): ${m.body_md}${attachmentNote(deps.db, m.id)}`; }).join("\n"); } diff --git a/src/main/services/agent-tools.ts b/src/main/services/agent-tools.ts index 3fdfb635..1d5c8d2c 100644 --- a/src/main/services/agent-tools.ts +++ b/src/main/services/agent-tools.ts @@ -10,6 +10,7 @@ import { grantConnectorTool, listConnectorsTool } from "./agent-connector-tools" import { bindSkillSlugs, grantSkillTool, listSkillsTool } from "./agent-skill-tools"; import { createDocumentTool, writeDeliverableTool } from "./agent-deliverable-tools"; import { companyStatusTool, GOAL_TOOLS, updateGoalTool } from "./agent-goal-tools"; +import { ATTACHMENT_TOOLS, readAttachmentTool } from "./agent-attachment-tools"; import { sendMessageTool } from "./agent-send-message"; import { withSavepoint } from "../database/connection"; import { optionalString, optionalStringArray, requireString } from "./agent-tool-args"; @@ -227,6 +228,7 @@ export const AGENT_TOOLS: ToolDefinition[] = [ }, }, ...GOAL_TOOLS, + ...ATTACHMENT_TOOLS, ]; export function dispatchAgentTool( @@ -257,6 +259,8 @@ export function dispatchAgentTool( return grantConnectorTool(deps, ctx, call.arguments, "grant"); case "revoke_connector": return grantConnectorTool(deps, ctx, call.arguments, "revoke"); + case "read_attachment": + return readAttachmentTool(deps, ctx, call.arguments); case "company_status": return companyStatusTool(deps, ctx); case "update_goal": diff --git a/src/main/services/chat-wake-context.ts b/src/main/services/chat-wake-context.ts index 3447d634..fd865b77 100644 --- a/src/main/services/chat-wake-context.ts +++ b/src/main/services/chat-wake-context.ts @@ -1,10 +1,11 @@ import type { DatabaseSync } from "node:sqlite"; import type { ChatContextBlock } from "./task-prompts"; +import { attachmentNote } from "./agent-attachment-tools"; const HISTORY_LIMIT = 20; const MESSAGE_SELECT = ` - select m.body_md, m.created_at, m.edited_at, m.sender_type, m.sender_id, + select m.id, m.body_md, m.created_at, m.edited_at, m.sender_type, m.sender_id, coalesce(a.name, e.display_name, m.sender_type) as sender_name, coalesce(a.role, e.role) as sender_role from chat_messages m @@ -12,6 +13,7 @@ const MESSAGE_SELECT = ` left join employees e on e.id = m.sender_id and m.sender_type = 'employee'`; interface ThreadMessageRow { + id: string; body_md: string; created_at: string; edited_at: string | null; @@ -46,15 +48,15 @@ export function loadChatWakeContext(db: DatabaseSync, orgId: string, agentId: st return { threadId: scope.sourceThreadId, spaceId: thread.space_id, - history: [...(omitted > 0 ? [`(${omitted} earlier messages in this thread are not shown.)`] : []), ...history.map(formatMessage)].join("\n"), - fresh: fresh.map(formatMessage).join("\n"), + history: [...(omitted > 0 ? [`(${omitted} earlier messages in this thread are not shown.)`] : []), ...history.map((m) => formatMessage(db, m))].join("\n"), + fresh: fresh.map((m) => formatMessage(db, m)).join("\n"), }; } -function formatMessage(m: ThreadMessageRow): string { +function formatMessage(db: DatabaseSync, m: ThreadMessageRow): string { const who = m.sender_role ? `${m.sender_name} (${m.sender_type}, ${m.sender_role})` : `${m.sender_name} (${m.sender_type})`; const edited = m.edited_at ? ` (edited ${utcMinute(m.edited_at)})` : ""; - return `[${utcMinute(m.created_at)}] ${who}${edited}: ${m.body_md}`; + return `[${utcMinute(m.created_at)}] ${who}${edited}: ${m.body_md}${attachmentNote(db, m.id)}`; } function utcMinute(iso: string): string { @@ -80,6 +82,6 @@ export function createRunMessagePoller( if (rows.length === 0) return null; watermark = rows[rows.length - 1].created_at; console.info(`[run-inbox] agent ${input.agentId}: ${rows.length} new message(s) injected mid-run`); - return ["Messages that arrived while you were working (they may change what you should do):", ...rows.map(formatMessage)].join("\n"); + return ["Messages that arrived while you were working (they may change what you should do):", ...rows.map((m) => formatMessage(db, m))].join("\n"); }; } diff --git a/tests/integration/chat-context.test.ts b/tests/integration/chat-context.test.ts index ab1dddb3..9e86c4a3 100644 --- a/tests/integration/chat-context.test.ts +++ b/tests/integration/chat-context.test.ts @@ -1,4 +1,5 @@ import { randomUUID } from "node:crypto"; +import { writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { describe, expect, test } from "vitest"; @@ -8,6 +9,8 @@ import { ProviderService } from "../../src/main/services/provider-service"; import { providerAdapters } from "../../src/main/services/provider-adapters"; import { TaskExecutionService } from "../../src/main/services/task-execution-service"; import { ChatService } from "../../src/main/services/chat-service"; +import { OrganizationService } from "../../src/main/services/organization-service"; +import { dispatchAgentTool, type AgentToolContext } from "../../src/main/services/agent-tools"; import { ensureCeoDirectSpace } from "../../src/main/services/ceo-direct-space"; import { claimNextQueuedRun } from "../../src/main/app/task-queue-worker"; import type { AgentTurnResult, ProviderToolTurnInput } from "../../src/main/services/provider-tool-turn"; @@ -51,8 +54,8 @@ async function founderLine() { const spaceId = ensureCeoDirectSpace(connection.db, orgId)!; const kickoff = claimNextQueuedRun(connection); await execution.executeRun(kickoff!.org_id, kickoff!.id); - const say = (bodyMd: string, threadId: string | null) => { - const id = chat.sendMessage({ orgId, spaceId, threadId, senderType: "employee", senderId: founderId, bodyMd, attachmentFileIds: [] }); + const say = (bodyMd: string, threadId: string | null, attachmentFileIds: string[] = []) => { + const id = chat.sendMessage({ orgId, spaceId, threadId, senderType: "employee", senderId: founderId, bodyMd, attachmentFileIds }); return (connection.db.prepare("select thread_id from chat_messages where id = ?").get(id) as { thread_id: string }).thread_id; }; const runWake = async () => { @@ -77,7 +80,17 @@ async function founderLine() { const row = connection.db.prepare("select id from chat_messages where body_md = ?").get(bodyMd) as { id: string }; chat.deleteMessage({ orgId, messageId: row.id, actorType: "employee", actorId: founderId }); }; - return { connection, wakePrompts, wakeInputs, say, runWake, runManually, edit, remove, onWakeTurn }; + const addFile = (name: string, content: string) => { + const path = join(tmpdir(), `ac-attach-${randomUUID()}-${name}`); + writeFileSync(path, content); + const fileId = `fil_${randomUUID()}`; + connection.db.prepare(` + insert into file_objects(id, org_id, name, kind, size_bytes, sha256, storage_path, owner_type, owner_id, created_at) + values (?, ?, ?, 'md', ?, 'x', ?, 'employee', ?, ?) + `).run(fileId, orgId, name, content.length, path, founderId, new Date().toISOString()); + return fileId; + }; + return { connection, orgId, ceoId, wakePrompts, wakeInputs, say, runWake, runManually, edit, remove, onWakeTurn, addFile }; } describe("chat wake context", () => { @@ -257,4 +270,21 @@ describe("chat wake context", () => { expect(h.connection.db.prepare("select count(*) as n from tasks where title like '%99-1234%' or description like '%99-1234%'").get()).toEqual({ n: 0 }); h.connection.close(); }); + + test("attachments travel with the message and can be read by file_id", async () => { + const h = await founderLine(); + const fileId = h.addFile("spec.md", "# Spec\nShip the booking API by Friday."); + h.say("Please review the attached spec.", null, [fileId]); + await h.runWake(); + expect(h.wakePrompts[0]).toContain(`Please review the attached spec. [attachments: spec.md (file_id=${fileId})]`); + + const deps = { db: h.connection.db, organization: new OrganizationService(h.connection.db), chat: new ChatService(h.connection.db) }; + const ctx: AgentToolContext = { + orgId: h.orgId, agent: { id: h.ceoId, modelProviderId: null, modelId: null }, + taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null, + }; + expect(dispatchAgentTool(deps, ctx, { id: "c1", name: "read_messages", arguments: { with: "Yang" } })).toContain(`[attachments: spec.md (file_id=${fileId})]`); + expect(dispatchAgentTool(deps, ctx, { id: "c2", name: "read_attachment", arguments: { file_id: fileId } })).toContain("Ship the booking API by Friday."); + h.connection.close(); + }); }); diff --git a/tests/unit/agent-attachment-tools.test.ts b/tests/unit/agent-attachment-tools.test.ts new file mode 100644 index 00000000..7b1a9c85 --- /dev/null +++ b/tests/unit/agent-attachment-tools.test.ts @@ -0,0 +1,62 @@ +import { randomUUID } from "node:crypto"; +import { writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, test } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OrganizationService } from "../../src/main/services/organization-service"; +import { ChatService } from "../../src/main/services/chat-service"; +import { dispatchAgentTool, type AgentToolContext } from "../../src/main/services/agent-tools"; + +function harness() { + const connection = new DatabaseConnection(join(tmpdir(), `ac-attach-${randomUUID()}.db`)); + const organization = new OrganizationService(connection.db); + let orgId = ""; + let otherOrgId = ""; + connection.transaction(() => { + orgId = organization.createOrganization({ organizationName: "A", employeeName: "Yang", mission: "m", goal: "g" }); + otherOrgId = organization.createOrganization({ organizationName: "B", employeeName: "Lee", mission: "m", goal: "g" }); + }); + const addFile = (owner: string, name: string, content: string | Buffer, scanState = "approved") => { + const path = join(tmpdir(), `ac-attach-${randomUUID()}-${name}`); + writeFileSync(path, content); + const fileId = `fil_${randomUUID()}`; + connection.db.prepare(` + insert into file_objects(id, org_id, name, kind, size_bytes, sha256, storage_path, owner_type, owner_id, created_at, scan_state) + values (?, ?, ?, 'md', 1, 'x', ?, 'employee', null, ?, ?) + `).run(fileId, owner, name, path, new Date().toISOString(), scanState); + return fileId; + }; + const ctx: AgentToolContext = { + orgId, agent: { id: "agt_x", modelProviderId: null, modelId: null }, + taskProjectId: null, taskGoalId: null, taskId: null, sourceThreadId: null, workspace: null, + }; + const read = (fileId: string) => dispatchAgentTool( + { db: connection.db, organization, chat: new ChatService(connection.db) }, + ctx, + { id: "c", name: "read_attachment", arguments: { file_id: fileId } }, + ); + return { connection, orgId, otherOrgId, addFile, read }; +} + +describe("read_attachment", () => { + test("reads text, caps long files, and names binary files", () => { + const h = harness(); + expect(h.read(h.addFile(h.orgId, "notes.md", "hello attachment"))).toBe("hello attachment"); + const long = h.read(h.addFile(h.orgId, "big.md", "x".repeat(30_000))); + expect(long).toContain("[truncated: showing the first 24000 of 30000 characters]"); + expect(h.read(h.addFile(h.orgId, "logo.png", Buffer.from([137, 80, 78, 71, 0, 1, 2])))).toBe("logo.png is a binary file (7 bytes); it has no text to read."); + expect(h.read(h.addFile(h.orgId, "blob.bin", Buffer.from([0x80, 0x81, 0xfe, 0xff])))).toBe("blob.bin is a binary file (4 bytes); it has no text to read."); + expect(h.read(h.addFile(h.orgId, "zh.md", "中文说明"))).toBe("中文说明"); + h.connection.close(); + }); + + test("refuses blocked files, unknown ids, and other companies' files", () => { + const h = harness(); + expect(h.read(h.addFile(h.orgId, "virus.txt", "bad", "blocked"))).toBe("ERROR: virus.txt was blocked by the file scan and cannot be read."); + expect(h.read("fil_missing")).toBe('ERROR: No file with file_id "fil_missing".'); + const foreign = h.addFile(h.otherOrgId, "secret.md", "other company"); + expect(h.read(foreign)).toBe(`ERROR: No file with file_id "${foreign}".`); + h.connection.close(); + }); +});