Agent 的 Message 表怎么设计：持久化、压缩与上下文

import { sql } from 'drizzle-orm'
import type { UIMessage } from 'ai'
import {
  bigint,
  check,
  foreignKey,
  index,
  integer,
  jsonb,
  pgEnum,
  pgTable,
  text,
  timestamp,
  unique,
  uniqueIndex,
  uuid,
} from 'drizzle-orm/pg-core'

type StoredMessageParts = UIMessage['parts']
type StoredMessageMetadata = Record<string, unknown>

export const messageRole = pgEnum('agent_message_role', [
  'system',
  'user',
  'assistant',
])

export const messageStatus = pgEnum('agent_message_status', [
  'pending',
  'streaming',
  'completed',
  'failed',
  'cancelled',
])

export const contextState = pgEnum('agent_message_context_state', [
  'eligible',
  'excluded',
  'redacted',
])

export const runStatus = pgEnum('agent_run_status', [
  'queued',
  'running',
  'completed',
  'failed',
  'cancelled',
])

export const agentThreads = pgTable(
  'agent_thread',
  {
    id: uuid('id').defaultRandom().primaryKey(),
    tenantId: uuid('tenant_id').notNull(),
    nextSequence: bigint('next_sequence', { mode: 'number' })
      .notNull()
      .default(0),
    lastMessageAt: timestamp('last_message_at', { withTimezone: true }),
    createdAt: timestamp('created_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
  },
  (table) => [
    unique('agent_thread_id_tenant_unique').on(table.id, table.tenantId),
    index('agent_thread_tenant_recent_idx').on(
      table.tenantId,
      table.lastMessageAt.desc()
    ),
  ]
)

export const agentRuns = pgTable(
  'agent_run',
  {
    id: uuid('id').defaultRandom().primaryKey(),
    tenantId: uuid('tenant_id').notNull(),
    threadId: uuid('thread_id').notNull(),
    status: runStatus('status').notNull().default('queued'),
    provider: text('provider').notNull(),
    model: text('model').notNull(),
    inputTokens: integer('input_tokens'),
    outputTokens: integer('output_tokens'),
    startedAt: timestamp('started_at', { withTimezone: true }),
    finishedAt: timestamp('finished_at', { withTimezone: true }),
    createdAt: timestamp('created_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
  },
  (table) => [
    unique('agent_run_id_tenant_thread_unique').on(
      table.id,
      table.tenantId,
      table.threadId
    ),
    foreignKey({
      columns: [table.threadId, table.tenantId],
      foreignColumns: [agentThreads.id, agentThreads.tenantId],
      name: 'agent_run_thread_tenant_fk',
    }).onDelete('cascade'),
    index('agent_run_tenant_thread_created_idx').on(
      table.tenantId,
      table.threadId,
      table.createdAt.desc()
    ),
  ]
)

export const agentMessages = pgTable(
  'agent_message',
  {
    id: uuid('id').defaultRandom().primaryKey(),
    tenantId: uuid('tenant_id').notNull(),
    threadId: uuid('thread_id').notNull(),
    runId: uuid('run_id'),
    sequence: bigint('sequence', { mode: 'number' }).notNull(),
    clientRequestId: text('client_request_id'),
    role: messageRole('role').notNull(),
    status: messageStatus('status').notNull().default('pending'),
    parts: jsonb('parts')
      .$type<StoredMessageParts>()
      .notNull()
      .default(sql`'[]'::jsonb`),
    metadata: jsonb('metadata')
      .$type<StoredMessageMetadata>()
      .notNull()
      .default(sql`'{}'::jsonb`),
    contextState: contextState('context_state')
      .notNull()
      .default('eligible'),
    exclusionReason: text('exclusion_reason'),
    schemaVersion: integer('schema_version').notNull().default(1),
    contentVersion: integer('content_version').notNull().default(1),
    byteSize: integer('byte_size').notNull().default(0),
    createdAt: timestamp('created_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
    updatedAt: timestamp('updated_at', { withTimezone: true })
      .notNull()
      .defaultNow(),
    completedAt: timestamp('completed_at', { withTimezone: true }),
    redactedAt: timestamp('redacted_at', { withTimezone: true }),
    deletedAt: timestamp('deleted_at', { withTimezone: true }),
  },
  (table) => [
    foreignKey({
      columns: [table.threadId, table.tenantId],
      foreignColumns: [agentThreads.id, agentThreads.tenantId],
      name: 'agent_message_thread_tenant_fk',
    }).onDelete('cascade'),
    foreignKey({
      columns: [table.runId, table.tenantId, table.threadId],
      foreignColumns: [agentRuns.id, agentRuns.tenantId, agentRuns.threadId],
      name: 'agent_message_run_tenant_thread_fk',
    }).onDelete('restrict'),
    uniqueIndex('agent_message_thread_sequence_uidx').on(
      table.threadId,
      table.sequence
    ),
    uniqueIndex('agent_message_client_request_uidx')
      .on(table.tenantId, table.threadId, table.clientRequestId)
      .where(sql`${table.clientRequestId} is not null`),
    index('agent_message_tenant_thread_recent_idx')
      .on(table.tenantId, table.threadId, table.sequence.desc())
      .where(sql`${table.deletedAt} is null`),
    index('agent_message_run_idx')
      .on(table.runId)
      .where(sql`${table.runId} is not null`),
    check('agent_message_sequence_positive', sql`${table.sequence} > 0`),
    check(
      'agent_message_parts_is_array',
      sql`jsonb_typeof(${table.parts}) = 'array'`
    ),
    check(
      'agent_message_content_version_positive',
      sql`${table.contentVersion} > 0`
    ),
  ]
)

import { and, eq, sql } from 'drizzle-orm'

const [slot] = await tx
  .update(agentThreads)
  .set({
    nextSequence: sql`${agentThreads.nextSequence} + 1`,
  })
  .where(
    and(
      eq(agentThreads.id, threadId),
      eq(agentThreads.tenantId, tenantId)
    )
  )
  .returning({ sequence: agentThreads.nextSequence })

if (!slot) {
  throw new Error('Agent thread not found')
}

import {
  convertToModelMessages,
  createUIMessageStreamResponse,
  streamText,
  toUIMessageStream,
  validateUIMessages,
} from 'ai'

const storedMessages = await loadAuthorizedMessages({ tenantId, threadId })
const candidateMessages = [...storedMessages, incomingUserMessage]

const uiMessages = await validateUIMessages({
  messages: candidateMessages,
  metadataSchema,
  dataSchemas,
  tools,
})

const modelMessages = await convertToModelMessages(uiMessages, { tools })
const result = streamText({
  model,
  instructions: trustedServerInstructions,
  messages: modelMessages,
  tools,
})

// 需要“断开后继续生成”时显式消费；同时要提供真正的取消接口。
void result.consumeStream()

return createUIMessageStreamResponse({
  stream: toUIMessageStream({
    stream: result.stream,
    originalMessages: uiMessages,
    generateMessageId: () => crypto.randomUUID(),
    onEnd: async ({ responseMessage, isAborted, finishReason }) => {
      await finalizeAssistantMessage({
        finishReason,
        isAborted,
        message: responseMessage,
        runId,
      })
    },
  }),
})

鉴权与租户隔离
  → schema 验证与版本迁移
  → context_state / 隐私 / 可见性过滤
  → UIMessage 转 ModelMessage
  → 无损裁剪 reasoning、旧工具轨迹和空消息
  → 选择最近对话、固定事实、相关记忆与未完成任务
  → 必要时生成或复用摘要 artifact
  → 按原子 turn 做最终硬截断
  → 写入 context snapshot
  → 调用模型

可用输入预算
= 模型上下文上限
- 最大输出预留
- instructions 与工具定义
- 当前用户输入和附件
- 安全余量

import { pruneMessages, ToolLoopAgent, type ModelMessage } from 'ai'

const agent = new ToolLoopAgent({
  model,
  instructions: trustedServerInstructions,
  tools,
  prepareStep: async ({ messages, stepNumber }) => {
    const budget = await calculateStepBudget({ messages, model, stepNumber })

    if (budget.fits) {
      return {}
    }

    const prunedMessages = pruneMessages({
      messages,
      reasoning: 'all',
      toolCalls: 'before-last-3-messages',
      emptyMessages: 'remove',
    })

    const compactedMessages: ModelMessage[] = await compactWithArtifacts({
      budget,
      messages: prunedMessages,
    })

    await saveContextSnapshot({
      messages: compactedMessages,
      policyVersion: 'context-v3',
      stepNumber,
    })

    return { messages: compactedMessages }
  },
})

alter table agent_message enable row level security;

create policy agent_message_tenant_policy on agent_message
using (tenant_id = current_setting('app.tenant_id', true)::uuid)
with check (tenant_id = current_setting('app.tenant_id', true)::uuid);