Fix agent send attribution auth (#4522)

## Summary
- include both from_principal/from_agent plus the active wake write
token on runtime-originated sends
- allow server send attribution when from_agent is backed by a valid
active agent write token
- keep spoofed from_agent rejection and add coverage for the write-token
path

## Tests
- pnpm --filter @electric-ax/agents-mcp build
- pnpm --filter @electric-sql/client build
- pnpm --filter @electric-ax/agents-runtime build
- cd packages/agents-server && pnpm vitest run
test/electric-agents-routes.test.ts -t "send endpoint"
- cd packages/agents-runtime && pnpm vitest run
test/process-wake.test.ts -t "executes send via server API"
- pnpm --filter @electric-ax/agents-runtime typecheck
- pnpm --filter @electric-ax/agents-server typecheck

Author: KyleAMathews

.changeset/fix-agent-send-auth.md

@@ -0,0 +1,6 @@
+---
+"@electric-ax/agents-runtime": patch
+"@electric-ax/agents-server": patch
+---
+
+Fix runtime-originated agent send attribution by sending `from_principal`, `from_agent`, and the active wake write token, and accepting `from_agent` when backed by a valid agent write token.

packages/agents-runtime/src/process-wake.ts

@@ -1227,7 +1227,9 @@ export async function processWake(
           payload: send.payload,
           type: send.type,
           afterMs: send.afterMs,
+          fromPrincipal: notification.principal?.url,
           fromAgent: entityUrl,
+          writeToken,
         })
         .then(() => ({ sent: true as const, targetUrl: send.targetUrl }))
 

packages/agents-runtime/src/runtime-server-client.ts

@@ -86,7 +86,9 @@ export interface SendEntityMessageOptions {
   afterMs?: number
   mode?: `immediate` | `queued` | `paused` | `steer`
   position?: string
+  fromPrincipal?: string
   fromAgent?: string
+  writeToken?: string
 }
 
 export interface RegisterWakeOptions {
@@ -290,18 +292,30 @@ export function createRuntimeServerClient(
     afterMs,
     mode,
     position,
+    fromPrincipal,
     fromAgent,
+    writeToken,
   }: SendEntityMessageOptions): Promise<void> => {
     const body: Record<string, unknown> = { payload }
     if (type !== undefined) body.type = type
     if (afterMs !== undefined) body.afterMs = afterMs
     if (mode !== undefined) body.mode = mode
     if (position !== undefined) body.position = position
+    if (fromPrincipal !== undefined) body.from_principal = fromPrincipal
     if (fromAgent !== undefined) body.from_agent = fromAgent
 
+    const headers = new Headers({ 'content-type': `application/json` })
+    if (writeToken) {
+      applyTokenHeader(
+        headers,
+        config.writeTokenHeader ?? `authorization`,
+        writeToken
+      )
+    }
+
     const response = await request(`${entityRpcPath(targetUrl)}/send`, {
       method: `POST`,
-      headers: { 'content-type': `application/json` },
+      headers,
       body: JSON.stringify(body),
     })
 

packages/agents-runtime/test/process-wake.test.ts

@@ -1350,6 +1350,16 @@ describe(`processWake`, () => {
       },
     })
 
+    fetchMock.mockResolvedValueOnce(
+      new Response(
+        JSON.stringify({ ok: true, writeToken: `mock-write-token` }),
+        {
+          status: 200,
+          headers: { 'content-type': `application/json` },
+        }
+      )
+    )
+
     await processWake(
       makeNotification({
         principal: {
@@ -1368,10 +1378,14 @@ describe(`processWake`, () => {
     expect(String(sendUrl)).toContain(`target-entity-2/send`)
     const body = JSON.parse(sendOpts!.body as string) as Record<string, unknown>
     expect(body.from).toBeUndefined()
+    expect(body.from_principal).toBe(`http://localhost:3000/test-agent/agent-1`)
     expect(body.from_agent).toBe(`http://localhost:3000/test-agent/agent-1`)
     expect((sendOpts!.headers as Headers).get(`electric-principal`)).toBe(
       `entity:test-agent/agent-1`
     )
+    expect((sendOpts!.headers as Headers).get(`authorization`)).toBe(
+      `Bearer mock-write-token`
+    )
     expect((body.payload as Record<string, unknown>).action).toBe(`ping`)
   })
 

packages/agents-server/src/routing/entities-router.ts

@@ -187,6 +187,19 @@ function agentUrlPath(value: string): string {
   }
 }
 
+async function hasValidAgentWriteToken(
+  request: AgentsRouteRequest,
+  ctx: TenantContext,
+  fromAgent: string
+): Promise<boolean> {
+  const agentUrl = agentUrlPath(fromAgent)
+  const token = writeTokenFromRequest(request)
+  if (!token) return false
+  const agentEntity = await ctx.entityManager.registry.getEntity(agentUrl)
+  if (!agentEntity) return false
+  return ctx.entityManager.isValidWriteToken(agentEntity, token)
+}
+
 const inboxMessageBodySchema = Type.Object({
   payload: Type.Optional(Type.Unknown()),
   position: Type.Optional(Type.String()),
@@ -1138,7 +1151,14 @@ async function sendEntity(
   }
   if (parsed.from_agent !== undefined) {
     const principalAgentUrl = agentUrlForPrincipal(principal)
-    if (agentUrlPath(parsed.from_agent) !== principalAgentUrl) {
+    const fromAgentUrl = agentUrlPath(parsed.from_agent)
+    const matchesPrincipalAgent = fromAgentUrl === principalAgentUrl
+    const hasAgentWriteToken = await hasValidAgentWriteToken(
+      request,
+      ctx,
+      parsed.from_agent
+    )
+    if (!matchesPrincipalAgent && !hasAgentWriteToken) {
       return apiError(
         400,
         ErrCodeInvalidRequest,

packages/agents-server/test/electric-agents-routes.test.ts

@@ -8,9 +8,10 @@ function createRequest(
   method: string,
   path: string,
   body?: unknown,
-  rawBody = false
+  rawBody = false,
+  headers?: HeadersInit
 ): Request {
-  const init: RequestInit = { method }
+  const init: RequestInit = { method, headers }
   if (body !== undefined) {
     init.body = rawBody ? String(body) : JSON.stringify(body)
   }
@@ -28,10 +29,11 @@ async function routeResponse(
     id: `dev-local`,
     key: `system:dev-local`,
     url: `/principal/system:dev-local`,
-  }
+  },
+  headers?: HeadersInit
 ): Promise<Response> {
   const result = await globalRouter.fetch(
-    createRequest(method, path, body, rawBody),
+    createRequest(method, path, body, rawBody, headers),
     {
       service: `test`,
       entityManager: manager,
@@ -790,6 +792,61 @@ describe(`ElectricAgentsRoutes send endpoint`, () => {
     })
   })
 
+  it(`allows from_agent values with a valid active agent write token`, async () => {
+    const targetEntity = { url: `/chat/test` }
+    const agentEntity = { url: `/horton/current`, write_token: `entity-token` }
+    const manager = {
+      registry: {
+        getEntity: vi.fn(async (url: string) => {
+          if (url === `/chat/test`) return targetEntity
+          if (url === `/horton/current`) return agentEntity
+          return null
+        }),
+        getEntityType: vi.fn(),
+      },
+      isValidWriteToken: vi.fn(
+        (entity, token) => entity === agentEntity && token === `claim-token`
+      ),
+      ensurePrincipal: vi.fn().mockResolvedValue(undefined),
+      send: vi.fn().mockResolvedValue(undefined),
+    } as any
+
+    const response = await routeResponse(
+      manager,
+      `POST`,
+      `/_electric/entities/chat/test/send`,
+      {
+        payload: { text: `hi` },
+        from_principal: `/principal/system:dev-local`,
+        from_agent: `/horton/current`,
+      },
+      false,
+      {
+        kind: `system`,
+        id: `dev-local`,
+        key: `system:dev-local`,
+        url: `/principal/system:dev-local`,
+      },
+      { 'electric-claim-token': `claim-token` }
+    )
+
+    expect(response.status).toBe(204)
+    expect(manager.isValidWriteToken).toHaveBeenCalledWith(
+      agentEntity,
+      `claim-token`
+    )
+    expect(manager.send).toHaveBeenCalledWith(`/chat/test`, {
+      from: `/principal/system:dev-local`,
+      from_principal: `/principal/system:dev-local`,
+      from_agent: `/horton/current`,
+      payload: { text: `hi` },
+      key: undefined,
+      type: undefined,
+      mode: undefined,
+      position: undefined,
+    })
+  })
+
   it(`allows matching from_agent values for agent principals`, async () => {
     const manager = {
       registry: {