fix(tess): redact chat persistence and egress (#725)

apps/gateway/src/chat/chat.gateway-redaction.spec.ts

@@ -0,0 +1,170 @@
+import { describe, expect, it, vi } from 'vitest';
+import { ChatGateway } from './chat.gateway.js';
+
+const CONVERSATION_ID = 'conversation-1';
+const CANARY = 'sk_canary12345678';
+
+type GatewayInternals = {
+  clientSessions: Map<string, unknown>;
+  relayEvent(client: unknown, conversationId: string, event: unknown): void;
+};
+
+function buildGateway() {
+  const brain = {
+    conversations: {
+      addMessage: vi.fn().mockResolvedValue(undefined),
+    },
+  };
+  const agentService = {
+    getSession: vi.fn().mockReturnValue(undefined),
+  };
+  const gateway = new ChatGateway(
+    agentService as never,
+    {} as never,
+    brain as never,
+    {} as never,
+    {} as never,
+    {} as never,
+  );
+
+  return { gateway: gateway as unknown as GatewayInternals, brain };
+}
+
+describe('ChatGateway redaction boundary', (): void => {
+  it('redacts a secret split across assistant deltas before egress and persistence', (): void => {
+    const { gateway } = buildGateway();
+    const client = {
+      connected: true,
+      id: 'client-1',
+      data: { user: { id: 'user-1' } },
+      emit: vi.fn(),
+    };
+    const session = {
+      conversationId: CONVERSATION_ID,
+      cleanup: vi.fn(),
+      assistantText: '',
+      toolCalls: [],
+      pendingToolCalls: new Map(),
+      scope: { userId: 'user-1', tenantId: 'tenant-1' },
+    };
+    gateway.clientSessions.set(client.id, session);
+
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: 'sk_canary' },
+    });
+
+    expect(JSON.stringify(client.emit.mock.calls)).not.toContain('sk_canary');
+
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: '12345678 ' },
+    });
+
+    expect(client.emit).toHaveBeenCalledWith('agent:text', {
+      conversationId: CONVERSATION_ID,
+      text: '[REDACTED_SECRET] ',
+    });
+    expect(session.assistantText).toBe(`${CANARY} `);
+    expect(JSON.stringify(client.emit.mock.calls)).not.toContain(CANARY);
+  });
+
+  it('retains a split secret label until its value can be redacted', (): void => {
+    const { gateway } = buildGateway();
+    const client = {
+      connected: true,
+      id: 'client-1',
+      data: { user: { id: 'user-1' } },
+      emit: vi.fn(),
+    };
+
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: 'token ' },
+    });
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: '=canaryvalue123 ' },
+    });
+
+    expect(client.emit).toHaveBeenCalledWith('agent:text', {
+      conversationId: CONVERSATION_ID,
+      text: '[REDACTED_SECRET] ',
+    });
+    expect(JSON.stringify(client.emit.mock.calls)).not.toContain('canaryvalue123');
+  });
+
+  it('holds a streamed private key until it can be redacted', (): void => {
+    const { gateway } = buildGateway();
+    const client = {
+      connected: true,
+      id: 'client-1',
+      data: { user: { id: 'user-1' } },
+      emit: vi.fn(),
+    };
+
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: '-----BEGIN PRIVATE KEY-----\ncanary' },
+    });
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: '\n-----END PRIVATE KEY-----' },
+    });
+
+    expect(client.emit).toHaveBeenCalledWith('agent:text', {
+      conversationId: CONVERSATION_ID,
+      text: '[REDACTED_SECRET]',
+    });
+    expect(JSON.stringify(client.emit.mock.calls)).not.toContain('canary');
+  });
+
+  it('drops an oversized unterminated stream fragment rather than retaining it', (): void => {
+    const { gateway } = buildGateway();
+    const client = {
+      connected: true,
+      id: 'client-1',
+      data: { user: { id: 'user-1' } },
+      emit: vi.fn(),
+    };
+
+    gateway.relayEvent(client, CONVERSATION_ID, {
+      type: 'message_update',
+      assistantMessageEvent: { type: 'text_delta', delta: 'x'.repeat(8_193) },
+    });
+
+    expect(client.emit).toHaveBeenCalledWith('agent:text', {
+      conversationId: CONVERSATION_ID,
+      text: '[REDACTED_STREAM_OVERFLOW]',
+    });
+  });
+
+  it('persists only redacted assistant content with classifications', (): void => {
+    const { gateway, brain } = buildGateway();
+    const client = {
+      connected: true,
+      id: 'client-1',
+      data: { user: { id: 'user-1' } },
+      emit: vi.fn(),
+    };
+    gateway.clientSessions.set(client.id, {
+      conversationId: CONVERSATION_ID,
+      cleanup: vi.fn(),
+      assistantText: CANARY,
+      toolCalls: [],
+      pendingToolCalls: new Map(),
+      scope: { userId: 'user-1', tenantId: 'tenant-1' },
+    });
+
+    gateway.relayEvent(client, CONVERSATION_ID, { type: 'agent_end' });
+
+    expect(brain.conversations.addMessage).toHaveBeenCalledWith(
+      expect.objectContaining({
+        content: '[REDACTED_SECRET]',
+        metadata: expect.objectContaining({ classifications: ['secret'] }),
+      }),
+      'user-1',
+    );
+    expect(JSON.stringify(brain.conversations.addMessage.mock.calls)).not.toContain(CANARY);
+  });
+});

apps/gateway/src/chat/chat.gateway.ts

@@ -18,6 +18,7 @@ import {
 } from '@mosaicstack/discord-plugin';
 import type { Auth } from '@mosaicstack/auth';
 import type { Brain } from '@mosaicstack/brain';
+import { redactSensitiveContent } from '@mosaicstack/log';
 import type {
   SetThinkingPayload,
   SlashCommandApprovalResultPayload,
@@ -63,6 +64,7 @@ interface ClientSession {
  * Keyed by conversationId, value is the model name to use.
  */
 const modelOverrides = new Map<string, string>();
+const MAX_REDACTION_BUFFER_LENGTH = 8_192;
 
 function isDiscordIngressEnvelope(value: unknown): value is DiscordIngressEnvelope {
   if (typeof value !== 'object' || value === null) return false;
@@ -107,6 +109,10 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
 
   private readonly logger = new Logger(ChatGateway.name);
   private readonly clientSessions = new Map<string, ClientSession>();
+  /** Raw stream fragments are kept in memory only until they are safe to redact and emit. */
+  private readonly textEgressBuffers = new Map<string, string>();
+  private readonly thinkingEgressBuffers = new Map<string, string>();
+  private readonly overflowedEgress = new Set<string>();
   private readonly discordReplayProtector = new DiscordReplayProtector();
 
   constructor(
@@ -155,6 +161,10 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
       );
       this.clientSessions.delete(client.id);
     }
+    this.textEgressBuffers.delete(client.id);
+    this.thinkingEgressBuffers.delete(client.id);
+    this.overflowedEgress.delete(this.egressKey(client, 'agent:text'));
+    this.overflowedEgress.delete(this.egressKey(client, 'agent:thinking'));
   }
 
   private getClientScope(client: Socket): ActorTenantScope | null {
@@ -317,7 +327,7 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
           {
             conversationId,
             role: 'user',
-            content: data.content,
+            content: redactSensitiveContent(data.content).content,
             metadata: {
               timestamp: new Date().toISOString(),
               ...(correlationId
@@ -327,6 +337,7 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
                     discordUserId: discordIngress?.userId,
                   }
                 : {}),
+              classifications: redactSensitiveContent(data.content).classifications,
             },
           },
           userId,
@@ -744,6 +755,127 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
     }
   }
 
+  private appendAndFlushRedactedEgress(
+    client: Socket,
+    conversationId: string,
+    eventName: 'agent:text' | 'agent:thinking',
+    buffers: Map<string, string>,
+    delta: string,
+  ): void {
+    const key = this.egressKey(client, eventName);
+    if (this.overflowedEgress.has(key)) return;
+
+    const buffered = `${buffers.get(client.id) ?? ''}${delta}`;
+    if (buffered.length > MAX_REDACTION_BUFFER_LENGTH) {
+      buffers.delete(client.id);
+      this.overflowedEgress.add(key);
+      client.emit(eventName, { conversationId, text: '[REDACTED_STREAM_OVERFLOW]' });
+      return;
+    }
+
+    buffers.set(client.id, buffered);
+    this.flushRedactedEgress(client, conversationId, eventName, buffers, false);
+  }
+
+  /**
+   * Holds any suffix that could become a secret, email, or phone number after a
+   * later stream chunk. This avoids relying on downstream redaction after data
+   * has already reached the socket.
+   */
+  private flushRedactedEgress(
+    client: Socket,
+    conversationId: string,
+    eventName: 'agent:text' | 'agent:thinking',
+    buffers: Map<string, string>,
+    final: boolean,
+  ): void {
+    const key = this.egressKey(client, eventName);
+    if (this.overflowedEgress.has(key)) {
+      if (final) this.overflowedEgress.delete(key);
+      return;
+    }
+
+    const buffered = buffers.get(client.id) ?? '';
+    const releaseLength = final ? buffered.length : this.safeRedactionPrefixLength(buffered);
+    const released = buffered.slice(0, releaseLength);
+    const pending = buffered.slice(releaseLength);
+
+    if (pending) {
+      buffers.set(client.id, pending);
+    } else {
+      buffers.delete(client.id);
+    }
+
+    if (released) {
+      client.emit(eventName, {
+        conversationId,
+        text: redactSensitiveContent(released).content,
+      });
+    }
+  }
+
+  private safeRedactionPrefixLength(content: string): number {
+    let retainedFrom = content.length;
+
+    // Retain the current token because it may become a split secret or email.
+    const token = /(?:^|\s)(\S*)$/.exec(content);
+    if (token) {
+      const matched = token[0] ?? '';
+      const trailingToken = token[1] ?? '';
+      retainedFrom = token.index + matched.length - trailingToken.length;
+    }
+
+    // The secret classifier accepts whitespace around ':' and '=', so preserve
+    // a pending label until its value and delimiter are both complete.
+    const pendingSecretLabel =
+      /(?:^|[^A-Za-z0-9_])((?:api[_-]?key|token|password|secret|bearer|authorization)\s*)$/i.exec(
+        content,
+      );
+    if (pendingSecretLabel) {
+      const label = pendingSecretLabel[1] ?? '';
+      retainedFrom = Math.min(
+        retainedFrom,
+        pendingSecretLabel.index + pendingSecretLabel[0].length - label.length,
+      );
+    }
+
+    const secretLabel = /(?:api[_-]?key|token|password|secret|authorization)\s*[:=]\s*$/i.exec(
+      content,
+    );
+    if (secretLabel) {
+      retainedFrom = Math.min(retainedFrom, secretLabel.index);
+    }
+
+    // Phone numbers can contain whitespace and punctuation; preserve the full
+    // trailing numeric candidate until a non-phone character establishes a boundary.
+    const phone = /(?:^|[^A-Za-z0-9_])(\+?\d[\d(). -]*)$/.exec(content);
+    if (phone) {
+      const matched = phone[0] ?? '';
+      const trailingPhoneCandidate = phone[1] ?? '';
+      retainedFrom = Math.min(
+        retainedFrom,
+        phone.index + matched.length - trailingPhoneCandidate.length,
+      );
+    }
+
+    const privateKeyStart = content.lastIndexOf('-----BEGIN');
+    if (privateKeyStart >= 0) {
+      const privateKey = content.slice(privateKeyStart);
+      if (/-----END(?: [A-Z]+)* KEY-----/.test(privateKey)) {
+        // Release the complete block in one pass so the full-block classifier can redact it.
+        retainedFrom = content.length;
+      } else {
+        retainedFrom = Math.min(retainedFrom, privateKeyStart);
+      }
+    }
+
+    return retainedFrom;
+  }
+
+  private egressKey(client: Socket, eventName: 'agent:text' | 'agent:thinking'): string {
+    return `${client.id}:${eventName}`;
+  }
+
   private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
     if (!client.connected) {
       this.logger.warn(
@@ -761,6 +893,10 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
           cs.toolCalls = [];
           cs.pendingToolCalls.clear();
         }
+        this.textEgressBuffers.set(client.id, '');
+        this.thinkingEgressBuffers.set(client.id, '');
+        this.overflowedEgress.delete(this.egressKey(client, 'agent:text'));
+        this.overflowedEgress.delete(this.egressKey(client, 'agent:thinking'));
         client.emit('agent:start', { conversationId });
         break;
       }
@@ -789,6 +925,20 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
             }
           : undefined;
 
+        this.flushRedactedEgress(
+          client,
+          conversationId,
+          'agent:text',
+          this.textEgressBuffers,
+          true,
+        );
+        this.flushRedactedEgress(
+          client,
+          conversationId,
+          'agent:thinking',
+          this.thinkingEgressBuffers,
+          true,
+        );
         client.emit('agent:end', {
           conversationId,
           usage: usagePayload,
@@ -831,8 +981,11 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
               {
                 conversationId,
                 role: 'assistant',
-                content: cs.assistantText,
-                metadata,
+                content: redactSensitiveContent(cs.assistantText).content,
+                metadata: {
+                  ...metadata,
+                  classifications: redactSensitiveContent(cs.assistantText).classifications,
+                },
               },
               userId,
             )
@@ -854,20 +1007,26 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
       case 'message_update': {
         const assistantEvent = event.assistantMessageEvent;
         if (assistantEvent.type === 'text_delta') {
-          // Accumulate assistant text for persistence
+          // Keep raw stream material in memory only; persist and emit only redacted text.
           const cs = this.clientSessions.get(client.id);
           if (cs) {
             cs.assistantText += assistantEvent.delta;
           }
-          client.emit('agent:text', {
+          this.appendAndFlushRedactedEgress(
+            client,
             conversationId,
-            text: assistantEvent.delta,
-          });
+            'agent:text',
+            this.textEgressBuffers,
+            assistantEvent.delta,
+          );
         } else if (assistantEvent.type === 'thinking_delta') {
-          client.emit('agent:thinking', {
+          this.appendAndFlushRedactedEgress(
+            client,
             conversationId,
-            text: assistantEvent.delta,
-          });
+            'agent:thinking',
+            this.thinkingEgressBuffers,
+            assistantEvent.delta,
+          );
         }
         break;
       }