From 4f56868545e3ef07c4f8b1a513c0507a973b544c Mon Sep 17 00:00:00 2001 From: Arvin Xu Date: Thu, 9 Apr 2026 10:09:05 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix:=20allow=20templates=20to=20?= =?UTF-8?q?specify=20policyLoad=20so=20default=20docs=20are=20fully=20inje?= =?UTF-8?q?cted=20(#13672)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐛 fix: allow templates to specify policyLoad so default docs are fully injected All documents were hardcoded to PolicyLoad.PROGRESSIVE on creation, causing CLAW template docs (IDENTITY, SOUL, BOOTSTRAP, AGENTS) to be progressively disclosed instead of fully injected into context. Co-Authored-By: Claude Opus 4.6 (1M context) * 🐛 fix: forward policyLoad through upsertDocument and persist on update - Add policyLoad to UpsertDocumentParams and pass it through to model - Add policyLoad param to update() so upsert's existing-document path writes the value instead of silently discarding it - Ensures re-running template init migrates pre-existing docs to ALWAYS Co-Authored-By: Claude Opus 4.6 (1M context) * ♻️ refactor: change update() to use named params object instead of positional args Co-Authored-By: Claude Opus 4.6 (1M context) * ♻️ refactor: change create() and upsert() to use named params object Co-Authored-By: Claude Opus 4.6 (1M context) * ✅ test: improve agentDocuments test coverage to 99% Add tests for uncovered branches: - normalizeLoadRule default branch (unknown rule) - explicit 'always' rule match - by-time-range with NaN dates - resolveDocumentLoadPosition fallback paths - composeToolPolicyUpdate with existing context values - upsert create path for new filenames - getAgentContext empty docs path Co-Authored-By: Claude Opus 4.6 (1M context) * 🐛 fix: preserve policyLoad when copying documents Co-Authored-By: Claude Opus 4.6 (1M context) * ✅ fix: align test assertion with refactored create() params object signature Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- packages/agent-templates/src/template.ts | 18 +- .../src/templates/claw/agent.ts | 3 +- .../src/templates/claw/bootstrap.ts | 3 +- .../src/templates/claw/identity.ts | 3 +- .../src/templates/claw/soul.ts | 3 +- .../__tests__/agentDocument.test.ts | 214 ++++++++---------- .../agentDocuments/__tests__/template.test.ts | 4 + .../models/agentDocuments/agentDocument.ts | 114 ++++++---- .../policy/__tests__/checks.test.ts | 26 +++ .../policy/__tests__/loadPolicy.test.ts | 29 +++ src/server/services/agentDocuments.test.ts | 11 +- src/server/services/agentDocuments.ts | 88 +++---- 12 files changed, 293 insertions(+), 223 deletions(-) diff --git a/packages/agent-templates/src/template.ts b/packages/agent-templates/src/template.ts index 93265c36a9..dc9460abdd 100644 --- a/packages/agent-templates/src/template.ts +++ b/packages/agent-templates/src/template.ts @@ -1,4 +1,9 @@ -import type { DocumentLoadFormat, DocumentLoadPosition, DocumentLoadRules } from './types'; +import type { + DocumentLoadFormat, + DocumentLoadPosition, + DocumentLoadRules, + PolicyLoad, +} from './types'; /** * Document Template Definition @@ -17,6 +22,8 @@ export interface DocumentTemplate { loadRules?: DocumentLoadRules; /** Additional metadata for the template */ metadata?: Record; + /** Controls whether this document is fully injected or progressively disclosed */ + policyLoad?: PolicyLoad; /** Default render format when the document is injected into context */ policyLoadFormat?: DocumentLoadFormat; /** Human-readable title for the template */ @@ -62,10 +69,11 @@ export class DocumentTemplateManager { options?: { description?: string; filename?: string; - policyLoadFormat?: DocumentLoadFormat; loadPosition?: DocumentLoadPosition; loadRules?: DocumentLoadRules; metadata?: Record; + policyLoad?: PolicyLoad; + policyLoadFormat?: DocumentLoadFormat; }, ): DocumentTemplate { return { @@ -73,10 +81,11 @@ export class DocumentTemplateManager { content, description: options?.description || `Template for ${title}`, filename: options?.filename || this.generateFilename(title), - policyLoadFormat: options?.policyLoadFormat, loadPosition: options?.loadPosition, loadRules: options?.loadRules, metadata: options?.metadata, + policyLoad: options?.policyLoad, + policyLoadFormat: options?.policyLoadFormat, }; } @@ -132,10 +141,11 @@ export class DocumentTemplateManager { options?: { description?: string; filename?: string; - policyLoadFormat?: DocumentLoadFormat; loadPosition?: DocumentLoadPosition; loadRules?: DocumentLoadRules; metadata?: Record; + policyLoad?: PolicyLoad; + policyLoadFormat?: DocumentLoadFormat; }, ): DocumentTemplate { const template = this.createBasic(title, content, options); diff --git a/packages/agent-templates/src/templates/claw/agent.ts b/packages/agent-templates/src/templates/claw/agent.ts index 35746f5b42..142b8a3b81 100644 --- a/packages/agent-templates/src/templates/claw/agent.ts +++ b/packages/agent-templates/src/templates/claw/agent.ts @@ -1,5 +1,5 @@ import type { DocumentTemplate } from '../../template'; -import { DocumentLoadFormat, DocumentLoadPosition } from '../../types'; +import { DocumentLoadFormat, DocumentLoadPosition, PolicyLoad } from '../../types'; import content from './AGENTS.md'; /** @@ -11,6 +11,7 @@ export const AGENT_DOCUMENT: DocumentTemplate = { title: 'Workspace', filename: 'AGENTS.md', description: 'How to use agent documents as durable state, working memory, and operating rules', + policyLoad: PolicyLoad.ALWAYS, policyLoadFormat: DocumentLoadFormat.FILE, loadPosition: DocumentLoadPosition.BEFORE_SYSTEM, loadRules: { diff --git a/packages/agent-templates/src/templates/claw/bootstrap.ts b/packages/agent-templates/src/templates/claw/bootstrap.ts index e5a5199ce9..e8f4073c92 100644 --- a/packages/agent-templates/src/templates/claw/bootstrap.ts +++ b/packages/agent-templates/src/templates/claw/bootstrap.ts @@ -1,5 +1,5 @@ import type { DocumentTemplate } from '../../template'; -import { DocumentLoadFormat, DocumentLoadPosition } from '../../types'; +import { DocumentLoadFormat, DocumentLoadPosition, PolicyLoad } from '../../types'; import content from './BOOTSTRAP.md'; /** @@ -13,6 +13,7 @@ export const BOOTSTRAP_DOCUMENT: DocumentTemplate = { title: 'Bootstrap', filename: 'BOOTSTRAP.md', description: 'First-run onboarding: discover identity, set up user profile, then self-destruct', + policyLoad: PolicyLoad.ALWAYS, policyLoadFormat: DocumentLoadFormat.FILE, loadPosition: DocumentLoadPosition.SYSTEM_APPEND, loadRules: { diff --git a/packages/agent-templates/src/templates/claw/identity.ts b/packages/agent-templates/src/templates/claw/identity.ts index e016e8487b..019810ab46 100644 --- a/packages/agent-templates/src/templates/claw/identity.ts +++ b/packages/agent-templates/src/templates/claw/identity.ts @@ -1,5 +1,5 @@ import type { DocumentTemplate } from '../../template'; -import { DocumentLoadFormat, DocumentLoadPosition } from '../../types'; +import { DocumentLoadFormat, DocumentLoadPosition, PolicyLoad } from '../../types'; import content from './IDENTITY.md'; /** @@ -11,6 +11,7 @@ export const IDENTITY_DOCUMENT: DocumentTemplate = { title: 'Identity', filename: 'IDENTITY.md', description: 'Name, creature type, vibe, and avatar identity', + policyLoad: PolicyLoad.ALWAYS, policyLoadFormat: DocumentLoadFormat.FILE, loadPosition: DocumentLoadPosition.SYSTEM_APPEND, loadRules: { diff --git a/packages/agent-templates/src/templates/claw/soul.ts b/packages/agent-templates/src/templates/claw/soul.ts index 49bee25025..df1a2fcf19 100644 --- a/packages/agent-templates/src/templates/claw/soul.ts +++ b/packages/agent-templates/src/templates/claw/soul.ts @@ -1,5 +1,5 @@ import type { DocumentTemplate } from '../../template'; -import { DocumentLoadFormat, DocumentLoadPosition } from '../../types'; +import { DocumentLoadFormat, DocumentLoadPosition, PolicyLoad } from '../../types'; import content from './SOUL.md'; /** @@ -12,6 +12,7 @@ export const SOUL_DOCUMENT: DocumentTemplate = { title: 'Soul', filename: 'SOUL.md', description: 'Core truths, boundaries, vibe, and continuity', + policyLoad: PolicyLoad.ALWAYS, policyLoadFormat: DocumentLoadFormat.FILE, loadPosition: DocumentLoadPosition.SYSTEM_APPEND, loadRules: { diff --git a/packages/database/src/models/agentDocuments/__tests__/agentDocument.test.ts b/packages/database/src/models/agentDocuments/__tests__/agentDocument.test.ts index 366bd6fae4..e6106a1310 100644 --- a/packages/database/src/models/agentDocuments/__tests__/agentDocument.test.ts +++ b/packages/database/src/models/agentDocuments/__tests__/agentDocument.test.ts @@ -41,15 +41,12 @@ beforeEach(async () => { describe('AgentDocumentModel', () => { describe('create', () => { it('should create an agent document with normalized policy and linked document row', async () => { - const result = await agentDocumentModel.create( - agentId, - 'identity.md', - 'line1\nline2', - DocumentLoadPosition.BEFORE_SYSTEM, - { maxTokens: 1024, priority: 2, rule: DocumentLoadRule.ALWAYS }, - 'claw', - { description: 'Identity policy', domain: 'ops' }, - ); + const result = await agentDocumentModel.create(agentId, 'identity.md', 'line1\nline2', { + loadPosition: DocumentLoadPosition.BEFORE_SYSTEM, + loadRules: { maxTokens: 1024, priority: 2, rule: DocumentLoadRule.ALWAYS }, + metadata: { description: 'Identity policy', domain: 'ops' }, + templateId: 'claw', + }); expect(result.agentId).toBe(agentId); expect(result.filename).toBe('identity.md'); @@ -108,24 +105,19 @@ describe('AgentDocumentModel', () => { describe('update and upsert', () => { it('should update content, metadata and policy projections', async () => { - const created = await agentDocumentModel.create( - agentId, - 'policy.md', - 'old', - DocumentLoadPosition.BEFORE_FIRST_USER, - { maxTokens: 100, priority: 8 }, - undefined, - { description: 'old desc', topic: 'old' }, - ); + const created = await agentDocumentModel.create(agentId, 'policy.md', 'old', { + loadPosition: DocumentLoadPosition.BEFORE_FIRST_USER, + loadRules: { maxTokens: 100, priority: 8 }, + metadata: { description: 'old desc', topic: 'old' }, + }); - await agentDocumentModel.update( - created.id, - 'new\ncontent', - DocumentLoadPosition.AFTER_KNOWLEDGE, - { maxTokens: 500, priority: 1 }, - { description: 'new desc', topic: 'new' }, - { context: { policyLoadFormat: DocumentLoadFormat.FILE } }, - ); + await agentDocumentModel.update(created.id, { + content: 'new\ncontent', + loadPosition: DocumentLoadPosition.AFTER_KNOWLEDGE, + loadRules: { maxTokens: 500, priority: 1 }, + metadata: { description: 'new desc', topic: 'new' }, + policy: { context: { policyLoadFormat: DocumentLoadFormat.FILE } }, + }); const updated = await agentDocumentModel.findById(created.id); expect(updated?.content).toBe('new\ncontent'); @@ -146,26 +138,30 @@ describe('AgentDocumentModel', () => { expect(updatedDoc?.description).toBe('new desc'); }); - it('should upsert by filename and merge metadata on updates', async () => { - const first = await agentDocumentModel.upsert( - agentId, - 'policy-upsert.md', - 'v1', - DocumentLoadPosition.BEFORE_FIRST_USER, - { priority: 9 }, - undefined, - { a: 1, description: 'v1' }, - ); + it('should upsert by creating a new document when filename does not exist', async () => { + const result = await agentDocumentModel.upsert(agentId, 'new-upsert.md', 'fresh', { + loadPosition: DocumentLoadPosition.BEFORE_SYSTEM, + loadRules: { priority: 5 }, + templateId: 'claw', + }); - const second = await agentDocumentModel.upsert( - agentId, - 'policy-upsert.md', - 'v2', - undefined, - { priority: 1, maxTokens: 900 }, - undefined, - { b: 2, description: 'v2' }, - ); + expect(result.filename).toBe('new-upsert.md'); + expect(result.content).toBe('fresh'); + expect(result.templateId).toBe('claw'); + expect(result.policy?.context?.position).toBe(DocumentLoadPosition.BEFORE_SYSTEM); + }); + + it('should upsert by filename and merge metadata on updates', async () => { + const first = await agentDocumentModel.upsert(agentId, 'policy-upsert.md', 'v1', { + loadPosition: DocumentLoadPosition.BEFORE_FIRST_USER, + loadRules: { priority: 9 }, + metadata: { a: 1, description: 'v1' }, + }); + + const second = await agentDocumentModel.upsert(agentId, 'policy-upsert.md', 'v2', { + loadRules: { priority: 1, maxTokens: 900 }, + metadata: { b: 2, description: 'v2' }, + }); expect(second.id).toBe(first.id); expect(second.content).toBe('v2'); @@ -193,15 +189,12 @@ describe('AgentDocumentModel', () => { }); it('should copy into a new record and keep policy/template metadata', async () => { - const created = await agentDocumentModel.create( - agentId, - 'copy-source.md', - 'copy me', - DocumentLoadPosition.BEFORE_SYSTEM, - { maxTokens: 200, priority: 3 }, - 'claw', - { description: 'source desc', domain: 'A' }, - ); + const created = await agentDocumentModel.create(agentId, 'copy-source.md', 'copy me', { + loadPosition: DocumentLoadPosition.BEFORE_SYSTEM, + loadRules: { maxTokens: 200, priority: 3 }, + metadata: { description: 'source desc', domain: 'A' }, + templateId: 'claw', + }); const copied = await agentDocumentModel.copy(created.id, 'Copied Title'); @@ -213,28 +206,38 @@ describe('AgentDocumentModel', () => { expect(copied?.policy?.context?.maxTokens).toBe(200); expect(copied?.metadata).toMatchObject({ description: 'source desc', domain: 'A' }); }); + + it('should preserve policyLoad when copying a document', async () => { + const created = await agentDocumentModel.create(agentId, 'always-doc.md', 'content', { + policyLoad: PolicyLoad.ALWAYS, + }); + + const copied = await agentDocumentModel.copy(created.id, 'Always Copy'); + + expect(copied?.policyLoad).toBe(PolicyLoad.ALWAYS); + }); }); describe('findByAgent and findByTemplate', () => { it('should return matched docs with parsed loadRules', async () => { - await agentDocumentModel.create(agentId, 'a.md', 'A', undefined, { - maxTokens: 100, - priority: 2, + await agentDocumentModel.create(agentId, 'a.md', 'A', { + loadRules: { maxTokens: 100, priority: 2 }, }); - await agentDocumentModel.create(agentId, 'b.md', 'B', undefined, { - maxTokens: 50, - priority: 1, + await agentDocumentModel.create(agentId, 'b.md', 'B', { + loadRules: { maxTokens: 50, priority: 1 }, + }); + await agentDocumentModel.create(agentId, 'c.md', 'C', { + loadRules: { priority: 9 }, + templateId: 'claw', + }); + await agentDocumentModel.create(agentId, 'd.md', 'D', { + loadRules: { priority: 8 }, + templateId: 'claw', + }); + await agentDocumentModel.create(secondAgentId, 'e.md', 'E', { + loadRules: { priority: 7 }, + templateId: 'claw', }); - await agentDocumentModel.create(agentId, 'c.md', 'C', undefined, { priority: 9 }, 'claw'); - await agentDocumentModel.create(agentId, 'd.md', 'D', undefined, { priority: 8 }, 'claw'); - await agentDocumentModel.create( - secondAgentId, - 'e.md', - 'E', - undefined, - { priority: 7 }, - 'claw', - ); const byAgent = await agentDocumentModel.findByAgent(agentId); expect(byAgent).toHaveLength(4); @@ -272,20 +275,14 @@ describe('AgentDocumentModel', () => { describe('updateToolLoadRule and loadable queries', () => { it('should apply tool load rule and exclude manual docs from loadable results', async () => { - const alwaysDoc = await agentDocumentModel.create( - agentId, - 'always.md', - 'always', - DocumentLoadPosition.BEFORE_FIRST_USER, - { priority: 2 }, - ); - const manualDoc = await agentDocumentModel.create( - agentId, - 'manual.md', - 'manual', - DocumentLoadPosition.BEFORE_FIRST_USER, - { priority: 1 }, - ); + const alwaysDoc = await agentDocumentModel.create(agentId, 'always.md', 'always', { + loadPosition: DocumentLoadPosition.BEFORE_FIRST_USER, + loadRules: { priority: 2 }, + }); + const manualDoc = await agentDocumentModel.create(agentId, 'manual.md', 'manual', { + loadPosition: DocumentLoadPosition.BEFORE_FIRST_USER, + loadRules: { priority: 1 }, + }); const updated = await agentDocumentModel.updateToolLoadRule(manualDoc.id, { keywordMatchMode: 'all', @@ -341,20 +338,14 @@ describe('AgentDocumentModel', () => { }); it('should group docs by position and sort by priority ascending', async () => { - await agentDocumentModel.create( - agentId, - 'p2.md', - 'p2', - DocumentLoadPosition.BEFORE_KNOWLEDGE, - { priority: 2 }, - ); - await agentDocumentModel.create( - agentId, - 'p1.md', - 'p1', - DocumentLoadPosition.BEFORE_KNOWLEDGE, - { priority: 1 }, - ); + await agentDocumentModel.create(agentId, 'p2.md', 'p2', { + loadPosition: DocumentLoadPosition.BEFORE_KNOWLEDGE, + loadRules: { priority: 2 }, + }); + await agentDocumentModel.create(agentId, 'p1.md', 'p1', { + loadPosition: DocumentLoadPosition.BEFORE_KNOWLEDGE, + loadRules: { priority: 1 }, + }); const grouped = await agentDocumentModel.getDocumentsByPosition(agentId); const docsAtPosition = grouped.get(DocumentLoadPosition.BEFORE_KNOWLEDGE) || []; @@ -393,23 +384,18 @@ describe('AgentDocumentModel', () => { expect(rawDoc).toBeDefined(); }); + it('should return empty string from getAgentContext when no loadable docs exist', async () => { + const context = await agentDocumentModel.getAgentContext(agentId); + expect(context).toBe(''); + }); + it('should soft delete by agent and by template', async () => { - const templateDoc = await agentDocumentModel.create( - agentId, - 'template-a.md', - 'A', - undefined, - undefined, - 'claw', - ); - const otherTemplateDoc = await agentDocumentModel.create( - agentId, - 'template-b.md', - 'B', - undefined, - undefined, - 'other', - ); + const templateDoc = await agentDocumentModel.create(agentId, 'template-a.md', 'A', { + templateId: 'claw', + }); + const otherTemplateDoc = await agentDocumentModel.create(agentId, 'template-b.md', 'B', { + templateId: 'other', + }); const secondAgentDoc = await agentDocumentModel.create(secondAgentId, 'agent-2.md', 'C'); await agentDocumentModel.deleteByTemplate(agentId, 'claw', 'template cleanup'); diff --git a/packages/database/src/models/agentDocuments/__tests__/template.test.ts b/packages/database/src/models/agentDocuments/__tests__/template.test.ts index e90d0c0626..d20f30e29e 100644 --- a/packages/database/src/models/agentDocuments/__tests__/template.test.ts +++ b/packages/database/src/models/agentDocuments/__tests__/template.test.ts @@ -69,6 +69,7 @@ describe('DocumentTemplateManager', () => { loadPosition: undefined, loadRules: undefined, metadata: undefined, + policyLoad: undefined, policyLoadFormat: undefined, title: 'Agent Notes', }); @@ -99,6 +100,7 @@ describe('DocumentTemplateManager', () => { loadPosition: DocumentLoadPosition.BEFORE_SYSTEM, loadRules, metadata: { scope: 'team' }, + policyLoad: undefined, policyLoadFormat: DocumentLoadFormat.FILE, title: 'Profile', }); @@ -153,6 +155,7 @@ describe('DocumentTemplateManager', () => { scope: 'private', variables: ['name'], }, + policyLoad: undefined, policyLoadFormat: undefined, title: 'Prompt', }); @@ -182,6 +185,7 @@ describe('DocumentTemplateManager', () => { scope: 'team', variables: ['name', 'team'], }, + policyLoad: undefined, policyLoadFormat: undefined, title: 'Cloned', }); diff --git a/packages/database/src/models/agentDocuments/agentDocument.ts b/packages/database/src/models/agentDocuments/agentDocument.ts index 227543818e..dc5f31c9a3 100644 --- a/packages/database/src/models/agentDocuments/agentDocument.ts +++ b/packages/database/src/models/agentDocuments/agentDocument.ts @@ -89,14 +89,28 @@ export class AgentDocumentModel { agentId: string, filename: string, content: string, - loadPosition?: DocumentLoadPosition, - loadRules?: DocumentLoadRules, - templateId?: string, - metadata?: Record, - policy?: AgentDocumentPolicy, - createdAt?: Date, - updatedAt?: Date, + params?: { + createdAt?: Date; + loadPosition?: DocumentLoadPosition; + loadRules?: DocumentLoadRules; + metadata?: Record; + policy?: AgentDocumentPolicy; + policyLoad?: PolicyLoad; + templateId?: string; + updatedAt?: Date; + }, ): Promise { + const { + createdAt, + loadPosition, + loadRules, + metadata, + policy, + policyLoad, + templateId, + updatedAt, + } = params ?? {}; + const title = filename.replace(/\.[^.]+$/, ''); const stats = this.getDocumentStats(content); const normalizedPolicy = normalizePolicy(loadPosition, loadRules, policy); @@ -131,7 +145,7 @@ export class AgentDocumentModel { accessShared: 0, agentId, createdAt, - policyLoad: PolicyLoad.PROGRESSIVE, + policyLoad: policyLoad ?? PolicyLoad.PROGRESSIVE, deleteReason: null, deletedAt: null, deletedByAgentId: null, @@ -155,12 +169,17 @@ export class AgentDocumentModel { async update( documentId: string, - content?: string, - loadPosition?: DocumentLoadPosition, - loadRules?: Partial, - metadata?: Record, - policy?: AgentDocumentPolicy, + params?: { + content?: string; + loadPosition?: DocumentLoadPosition; + loadRules?: Partial; + metadata?: Record; + policy?: AgentDocumentPolicy; + policyLoad?: PolicyLoad; + }, ): Promise { + const { content, loadPosition, loadRules, metadata, policy, policyLoad } = params ?? {}; + const existing = await this.findById(documentId); if (!existing) return; @@ -191,6 +210,7 @@ export class AgentDocumentModel { policyLoadFormat: mergedPolicy.context?.policyLoadFormat || DocumentLoadFormat.RAW, policyLoadPosition: mergedPolicy.context?.position || DocumentLoadPosition.BEFORE_FIRST_USER, policyLoadRule: mergedPolicy.context?.rule || DocumentLoadRule.ALWAYS, + ...(policyLoad !== undefined && { policyLoad }), }; await this.db.transaction(async (trx) => { @@ -253,17 +273,16 @@ export class AgentDocumentModel { ? buildDocumentFilename(title, existing.filename) : `copy-${Date.now()}-${existing.filename}`; - return this.create( - existing.agentId, - filename, - existing.content, - (existing.policy?.context?.position as DocumentLoadPosition | undefined) || + return this.create(existing.agentId, filename, existing.content, { + loadPosition: + (existing.policy?.context?.position as DocumentLoadPosition | undefined) || DocumentLoadPosition.BEFORE_FIRST_USER, - parseLoadRules(existing), - existing.templateId || undefined, - existing.metadata || undefined, - existing.policy || undefined, - ); + loadRules: parseLoadRules(existing), + metadata: existing.metadata || undefined, + policy: existing.policy || undefined, + policyLoad: existing.policyLoad as PolicyLoad | undefined, + templateId: existing.templateId || undefined, + }); } async updateToolLoadRule( @@ -316,14 +335,28 @@ export class AgentDocumentModel { agentId: string, filename: string, content: string, - loadPosition?: DocumentLoadPosition, - loadRules?: DocumentLoadRules, - templateId?: string, - metadata?: Record, - policy?: AgentDocumentPolicy, - createdAt?: Date, - updatedAt?: Date, + params?: { + createdAt?: Date; + loadPosition?: DocumentLoadPosition; + loadRules?: DocumentLoadRules; + metadata?: Record; + policy?: AgentDocumentPolicy; + policyLoad?: PolicyLoad; + templateId?: string; + updatedAt?: Date; + }, ): Promise { + const { + createdAt, + loadPosition, + loadRules, + metadata, + policy, + policyLoad, + templateId, + updatedAt, + } = params ?? {}; + const existing = await this.findByFilename(agentId, filename); if (existing) { @@ -333,23 +366,28 @@ export class AgentDocumentModel { ? { ...existing.metadata, ...metadata } : (existing.metadata ?? undefined); - await this.update(existing.id, content, loadPosition, mergedRules, mergedMetadata, policy); + await this.update(existing.id, { + content, + loadPosition, + loadRules: mergedRules, + metadata: mergedMetadata, + policy, + policyLoad, + }); return (await this.findByFilename(agentId, filename))!; } - return this.create( - agentId, - filename, - content, + return this.create(agentId, filename, content, { + createdAt, loadPosition, loadRules, - templateId, metadata, policy, - createdAt, + policyLoad, + templateId, updatedAt, - ); + }); } async findByAgent(agentId: string): Promise { diff --git a/packages/database/src/models/agentDocuments/policy/__tests__/checks.test.ts b/packages/database/src/models/agentDocuments/policy/__tests__/checks.test.ts index c0855ffc4b..ec4b247629 100644 --- a/packages/database/src/models/agentDocuments/policy/__tests__/checks.test.ts +++ b/packages/database/src/models/agentDocuments/policy/__tests__/checks.test.ts @@ -54,6 +54,32 @@ describe('agentDocuments checks', () => { expect(composed.policy.context?.keywords).toEqual(['risk']); }); + it('resolves document position from policyLoadPosition fallback', () => { + expect( + resolveDocumentLoadPosition({ + policy: { context: {} }, + policyLoadPosition: DocumentLoadPosition.AFTER_KNOWLEDGE, + }), + ).toBe(DocumentLoadPosition.AFTER_KNOWLEDGE); + + expect( + resolveDocumentLoadPosition({ + policy: null, + policyLoadPosition: undefined as any, + }), + ).toBe(DocumentLoadPosition.BEFORE_FIRST_USER); + }); + + it('composes tool policy with rule/format from existing context when not in rule', () => { + const composed = composeToolPolicyUpdate( + { context: { policyLoadFormat: DocumentLoadFormat.FILE, rule: DocumentLoadRule.BY_REGEXP } }, + {}, + ); + + expect(composed.policyLoadFormat).toBe(DocumentLoadFormat.FILE); + expect(composed.policyLoadRule).toBe(DocumentLoadRule.BY_REGEXP); + }); + it('parses load rules and resolves document position', () => { const doc = { policy: { diff --git a/packages/database/src/models/agentDocuments/policy/__tests__/loadPolicy.test.ts b/packages/database/src/models/agentDocuments/policy/__tests__/loadPolicy.test.ts index 6c125b4566..063b74f378 100644 --- a/packages/database/src/models/agentDocuments/policy/__tests__/loadPolicy.test.ts +++ b/packages/database/src/models/agentDocuments/policy/__tests__/loadPolicy.test.ts @@ -62,6 +62,35 @@ describe('agentDocuments load policy checks', () => { ).toBe(true); }); + it('returns true for unknown load rule (normalizeLoadRule default branch)', () => { + expect( + matchesLoadRules( + { loadRules: { rule: 'unknown-rule' as any } }, + { currentUserMessage: 'hello' }, + ), + ).toBe(true); + }); + + it('returns true for explicitly set always rule', () => { + expect( + matchesLoadRules({ loadRules: { rule: 'always' } }, { currentUserMessage: 'anything' }), + ).toBe(true); + }); + + it('rejects by-time-range with NaN dates', () => { + expect( + matchesLoadRules( + { + loadRules: { + rule: 'by-time-range', + timeRange: { from: 'not-a-date', to: 'also-not-a-date' }, + }, + }, + { currentTime: new Date() }, + ), + ).toBe(false); + }); + it('composes load-rule check through shouldInjectDocument', () => { const doc = { loadRules: { keywords: ['release'], rule: 'by-keywords' as const }, diff --git a/src/server/services/agentDocuments.test.ts b/src/server/services/agentDocuments.test.ts index 78fc317450..ef50e17c56 100644 --- a/src/server/services/agentDocuments.test.ts +++ b/src/server/services/agentDocuments.test.ts @@ -44,16 +44,7 @@ describe('AgentDocumentsService', () => { expect(mockModel.findByFilename).toHaveBeenNthCalledWith(1, 'agent-1', 'note.md'); expect(mockModel.findByFilename).toHaveBeenNthCalledWith(2, 'agent-1', 'note-2.md'); - expect(mockModel.create).toHaveBeenCalledWith( - 'agent-1', - 'note-2.md', - 'content', - undefined, - undefined, - undefined, - undefined, - undefined, - ); + expect(mockModel.create).toHaveBeenCalledWith('agent-1', 'note-2.md', 'content', undefined); expect(result).toEqual({ id: 'new-doc', filename: 'note-2.md' }); }); diff --git a/src/server/services/agentDocuments.ts b/src/server/services/agentDocuments.ts index 5e18b584ca..33c7446731 100644 --- a/src/server/services/agentDocuments.ts +++ b/src/server/services/agentDocuments.ts @@ -5,6 +5,7 @@ import { type DocumentLoadRules, type DocumentTemplateSet, getDocumentTemplate, + type PolicyLoad, } from '@lobechat/agent-templates'; import type { LobeChatDatabase } from '@lobechat/database'; @@ -26,6 +27,7 @@ interface UpsertDocumentParams { loadRules?: DocumentLoadRules; metadata?: Record; policy?: AgentDocumentPolicy; + policyLoad?: PolicyLoad; templateId?: string; updatedAt?: Date; } @@ -45,11 +47,13 @@ export class AgentDocumentsService { agentId: string, title: string, content: string, - loadPosition?: DocumentLoadPosition, - loadRules?: DocumentLoadRules, - templateId?: string, - metadata?: Record, - policy?: AgentDocumentPolicy, + params?: { + loadPosition?: DocumentLoadPosition; + loadRules?: DocumentLoadRules; + metadata?: Record; + policy?: AgentDocumentPolicy; + templateId?: string; + }, ) { const baseFilename = buildDocumentFilename(title); const extensionMatch = baseFilename.match(/(\.[^./\\]+)$/); @@ -70,16 +74,7 @@ export class AgentDocumentsService { suffix += 1; } - return this.agentDocumentModel.create( - agentId, - filename, - content, - loadPosition, - loadRules, - templateId, - metadata, - policy, - ); + return this.agentDocumentModel.create(agentId, filename, content, params); } /** @@ -92,22 +87,16 @@ export class AgentDocumentsService { const templateSet = getDocumentTemplate(templateId); for (const template of templateSet.templates) { - await this.agentDocumentModel.upsert( - agentId, - template.filename, - template.content, - template.loadPosition, - template.loadRules, - templateId, - template.metadata, - template.policyLoadFormat - ? { - context: { - policyLoadFormat: template.policyLoadFormat, - }, - } + await this.agentDocumentModel.upsert(agentId, template.filename, template.content, { + loadPosition: template.loadPosition, + loadRules: template.loadRules, + metadata: template.metadata, + policy: template.policyLoadFormat + ? { context: { policyLoadFormat: template.policyLoadFormat } } : undefined, - ); + policyLoad: template.policyLoad, + templateId, + }); } } @@ -116,22 +105,16 @@ export class AgentDocumentsService { */ async initializeFromCustomTemplate(agentId: string, templateSet: DocumentTemplateSet) { for (const template of templateSet.templates) { - await this.agentDocumentModel.upsert( - agentId, - template.filename, - template.content, - template.loadPosition, - template.loadRules, - templateSet.id, - template.metadata, - template.policyLoadFormat - ? { - context: { - policyLoadFormat: template.policyLoadFormat, - }, - } + await this.agentDocumentModel.upsert(agentId, template.filename, template.content, { + loadPosition: template.loadPosition, + loadRules: template.loadRules, + metadata: template.metadata, + policy: template.policyLoadFormat + ? { context: { policyLoadFormat: template.policyLoadFormat } } : undefined, - ); + policyLoad: template.policyLoad, + templateId: templateSet.id, + }); } } @@ -209,21 +192,20 @@ export class AgentDocumentsService { templateId, metadata, policy, + policyLoad, createdAt, updatedAt, }: UpsertDocumentParams) { - return this.agentDocumentModel.upsert( - agentId, - filename, - content, + return this.agentDocumentModel.upsert(agentId, filename, content, { + createdAt, loadPosition, loadRules, - templateId, metadata, policy, - createdAt, + policyLoad, + templateId, updatedAt, - ); + }); } async createDocument(agentId: string, title: string, content: string) { @@ -357,7 +339,7 @@ export class AgentDocumentsService { const doc = await this.getDocumentByIdInAgent(documentId, expectedAgentId); if (!doc) return undefined; - await this.agentDocumentModel.update(documentId, content); + await this.agentDocumentModel.update(documentId, { content }); return this.agentDocumentModel.findById(documentId); }