[dev] [tofikwest] tofik/agent-chat-history-clears-on-close (#1964)

* feat(api): add assistant chat history management endpoints

* refactor(api): extract user context validation into a separate method

---------

Co-authored-by: Tofik Hasanov <[email protected]>

Author: github-actions[bot]

apps/api/package.json

@@ -28,6 +28,7 @@
     "@trigger.dev/sdk": "4.0.6",
     "@trycompai/db": "^1.3.20",
     "@trycompai/email": "workspace:*",
+    "@upstash/redis": "^1.34.2",
     "@upstash/vector": "^1.2.2",
     "adm-zip": "^0.5.16",
     "ai": "^5.0.60",

apps/api/src/app.module.ts

@@ -29,6 +29,7 @@ import { IntegrationPlatformModule } from './integration-platform/integration-pl
 import { CloudSecurityModule } from './cloud-security/cloud-security.module';
 import { BrowserbaseModule } from './browserbase/browserbase.module';
 import { TaskManagementModule } from './task-management/task-management.module';
+import { AssistantChatModule } from './assistant-chat/assistant-chat.module';
 
 @Module({
   imports: [
@@ -70,6 +71,7 @@ import { TaskManagementModule } from './task-management/task-management.module';
     CloudSecurityModule,
     BrowserbaseModule,
     TaskManagementModule,
+    AssistantChatModule,
   ],
   controllers: [AppController],
   providers: [

apps/api/src/assistant-chat/assistant-chat.controller.ts

@@ -0,0 +1,120 @@
+import {
+  BadRequestException,
+  Body,
+  Controller,
+  Delete,
+  Get,
+  Put,
+  UseGuards,
+} from '@nestjs/common';
+import {
+  ApiHeader,
+  ApiOperation,
+  ApiResponse,
+  ApiSecurity,
+  ApiTags,
+} from '@nestjs/swagger';
+import { AuthContext } from '../auth/auth-context.decorator';
+import { HybridAuthGuard } from '../auth/hybrid-auth.guard';
+import type { AuthContext as AuthContextType } from '../auth/types';
+import { SaveAssistantChatHistoryDto } from './assistant-chat.dto';
+import { AssistantChatService } from './assistant-chat.service';
+import type { AssistantChatMessage } from './assistant-chat.types';
+
+@ApiTags('Assistant Chat')
+@Controller({ path: 'assistant-chat', version: '1' })
+@UseGuards(HybridAuthGuard)
+@ApiSecurity('apikey')
+@ApiHeader({
+  name: 'X-Organization-Id',
+  description:
+    'Organization ID (required for JWT auth, optional for API key auth)',
+  required: false,
+})
+export class AssistantChatController {
+  constructor(private readonly assistantChatService: AssistantChatService) {}
+
+  private getUserScopedContext(auth: AuthContextType): { organizationId: string; userId: string } {
+    // Defensive checks (should already be guaranteed by HybridAuthGuard + AuthContext decorator)
+    if (!auth.organizationId) {
+      throw new BadRequestException('Organization ID is required');
+    }
+
+    if (auth.isApiKey) {
+      throw new BadRequestException(
+        'Assistant chat history is only available for user-authenticated requests (Bearer JWT).',
+      );
+    }
+
+    if (!auth.userId) {
+      throw new BadRequestException('User ID is required');
+    }
+
+    return { organizationId: auth.organizationId, userId: auth.userId };
+  }
+
+  @Get('history')
+  @ApiOperation({
+    summary: 'Get assistant chat history',
+    description:
+      'Returns the current user-scoped assistant chat history (ephemeral session context).',
+  })
+  @ApiResponse({
+    status: 200,
+    description: 'Chat history retrieved',
+    schema: {
+      type: 'object',
+      properties: {
+        messages: { type: 'array', items: { type: 'object' } },
+      },
+    },
+  })
+  async getHistory(@AuthContext() auth: AuthContextType): Promise<{ messages: AssistantChatMessage[] }> {
+    const { organizationId, userId } = this.getUserScopedContext(auth);
+
+    const messages = await this.assistantChatService.getHistory({
+      organizationId,
+      userId,
+    });
+
+    return { messages };
+  }
+
+  @Put('history')
+  @ApiOperation({
+    summary: 'Save assistant chat history',
+    description:
+      'Replaces the current user-scoped assistant chat history (ephemeral session context).',
+  })
+  async saveHistory(
+    @AuthContext() auth: AuthContextType,
+    @Body() dto: SaveAssistantChatHistoryDto,
+  ): Promise<{ success: true }> {
+    const { organizationId, userId } = this.getUserScopedContext(auth);
+
+    await this.assistantChatService.saveHistory(
+      { organizationId, userId },
+      dto.messages,
+    );
+
+    return { success: true };
+  }
+
+  @Delete('history')
+  @ApiOperation({
+    summary: 'Clear assistant chat history',
+    description: 'Deletes the current user-scoped assistant chat history.',
+  })
+  async clearHistory(@AuthContext() auth: AuthContextType): Promise<{ success: true }> {
+    const { organizationId, userId } = this.getUserScopedContext(auth);
+
+    await this.assistantChatService.clearHistory({
+      organizationId,
+      userId,
+    });
+
+    return { success: true };
+  }
+}
+
+

apps/api/src/assistant-chat/assistant-chat.dto.ts

@@ -0,0 +1,31 @@
+import { ApiProperty } from '@nestjs/swagger';
+import { IsArray, IsIn, IsNumber, IsString, ValidateNested } from 'class-validator';
+import { Type } from 'class-transformer';
+
+export class AssistantChatMessageDto {
+  @ApiProperty({ example: 'msg_abc123' })
+  @IsString()
+  id!: string;
+
+  @ApiProperty({ enum: ['user', 'assistant'], example: 'user' })
+  @IsIn(['user', 'assistant'])
+  role!: 'user' | 'assistant';
+
+  @ApiProperty({ example: 'How do I invite a teammate?' })
+  @IsString()
+  text!: string;
+
+  @ApiProperty({ example: 1735781554000, description: 'Unix epoch millis' })
+  @IsNumber()
+  createdAt!: number;
+}
+
+export class SaveAssistantChatHistoryDto {
+  @ApiProperty({ type: [AssistantChatMessageDto] })
+  @IsArray()
+  @ValidateNested({ each: true })
+  @Type(() => AssistantChatMessageDto)
+  messages!: AssistantChatMessageDto[];
+}
+
+

apps/api/src/assistant-chat/assistant-chat.module.ts

@@ -0,0 +1,13 @@
+import { Module } from '@nestjs/common';
+import { AuthModule } from '../auth/auth.module';
+import { AssistantChatController } from './assistant-chat.controller';
+import { AssistantChatService } from './assistant-chat.service';
+
+@Module({
+  imports: [AuthModule],
+  controllers: [AssistantChatController],
+  providers: [AssistantChatService],
+})
+export class AssistantChatModule {}
+
+

apps/api/src/assistant-chat/assistant-chat.service.ts

@@ -0,0 +1,53 @@
+import { Injectable } from '@nestjs/common';
+import { z } from 'zod';
+import { assistantChatRedisClient } from './upstash-redis.client';
+import type { AssistantChatMessage } from './assistant-chat.types';
+
+const StoredMessageSchema = z.object({
+  id: z.string(),
+  role: z.enum(['user', 'assistant']),
+  text: z.string(),
+  createdAt: z.number(),
+});
+
+const StoredMessagesSchema = z.array(StoredMessageSchema);
+
+type GetAssistantChatKeyParams = {
+  organizationId: string;
+  userId: string;
+};
+
+const getAssistantChatKey = ({ organizationId, userId }: GetAssistantChatKeyParams): string => {
+  return `assistant-chat:v1:${organizationId}:${userId}`;
+};
+
+@Injectable()
+export class AssistantChatService {
+  /**
+   * Default TTL is 7 days. This is intended to behave like "session context"
+   * rather than a long-term, searchable archive.
+   */
+  private readonly ttlSeconds = Number(process.env.ASSISTANT_CHAT_TTL_SECONDS ?? 60 * 60 * 24 * 7);
+
+  async getHistory(params: GetAssistantChatKeyParams): Promise<AssistantChatMessage[]> {
+    const key = getAssistantChatKey(params);
+    const raw = await assistantChatRedisClient.get<unknown>(key);
+    const parsed = StoredMessagesSchema.safeParse(raw);
+    if (!parsed.success) return [];
+    return parsed.data;
+  }
+
+  async saveHistory(params: GetAssistantChatKeyParams, messages: AssistantChatMessage[]): Promise<void> {
+    const key = getAssistantChatKey(params);
+    // Always validate before writing to keep the cache shape stable.
+    const validated = StoredMessagesSchema.parse(messages);
+    await assistantChatRedisClient.set(key, validated, { ex: this.ttlSeconds });
+  }
+
+  async clearHistory(params: GetAssistantChatKeyParams): Promise<void> {
+    const key = getAssistantChatKey(params);
+    await assistantChatRedisClient.del(key);
+  }
+}
+
+

apps/api/src/assistant-chat/assistant-chat.types.ts

@@ -0,0 +1,8 @@
+export type AssistantChatMessage = {
+  id: string;
+  role: 'user' | 'assistant';
+  text: string;
+  createdAt: number;
+};
+
+

apps/api/src/assistant-chat/upstash-redis.client.ts

@@ -0,0 +1,45 @@
+import { Redis } from '@upstash/redis';
+
+/**
+ * Upstash Redis client wrapper for the NestJS API.
+ *
+ * NOTE: We do NOT import `server-only` here because this code runs in Node (Nest),
+ * not Next.js Server Components.
+ */
+class InMemoryRedis {
+  private storage = new Map<string, { value: unknown; expiresAt?: number }>();
+
+  async get<T = unknown>(key: string): Promise<T | null> {
+    const record = this.storage.get(key);
+    if (!record) return null;
+    if (record.expiresAt && record.expiresAt <= Date.now()) {
+      this.storage.delete(key);
+      return null;
+    }
+    return record.value as T;
+  }
+
+  async set(key: string, value: unknown, options?: { ex?: number }): Promise<'OK'> {
+    const expiresAt = options?.ex ? Date.now() + options.ex * 1000 : undefined;
+    this.storage.set(key, { value, expiresAt });
+    return 'OK';
+  }
+
+  async del(key: string): Promise<number> {
+    const existed = this.storage.delete(key);
+    return existed ? 1 : 0;
+  }
+}
+
+const hasUpstashConfig =
+  !!process.env.UPSTASH_REDIS_REST_URL && !!process.env.UPSTASH_REDIS_REST_TOKEN;
+
+export const assistantChatRedisClient: Pick<Redis, 'get' | 'set' | 'del'> =
+  hasUpstashConfig
+    ? new Redis({
+        url: process.env.UPSTASH_REDIS_REST_URL!,
+        token: process.env.UPSTASH_REDIS_REST_TOKEN!,
+      })
+    : (new InMemoryRedis() as unknown as Pick<Redis, 'get' | 'set' | 'del'>);
+
+

apps/app/src/app/api/chat/route.ts

@@ -2,6 +2,7 @@ import { tools } from '@/data/tools';
 import { env } from '@/env.mjs';
 import { auth } from '@/utils/auth';
 import { openai } from '@ai-sdk/openai';
+import { db } from '@db';
 import { type UIMessage, convertToModelMessages, streamText } from 'ai';
 import { headers } from 'next/headers';
 import { NextResponse } from 'next/server';
@@ -19,10 +20,39 @@ export async function POST(req: Request) {
     headers: await headers(),
   });
 
-  if (!session?.session.activeOrganizationId) {
+  if (!session?.user) {
     return new Response('Unauthorized', { status: 401 });
   }
 
+  const organizationIdFromHeader = req.headers.get('x-organization-id')?.trim();
+  const organizationIdFromSession = session.session.activeOrganizationId;
+
+  // Prefer deterministic org context from URL → client header.
+  const organizationId = organizationIdFromHeader ?? organizationIdFromSession;
+
+  if (!organizationId) {
+    return NextResponse.json(
+      { error: 'Organization context required (missing X-Organization-Id).' },
+      { status: 400 },
+    );
+  }
+
+  // Validate the user is a member of the requested org.
+  const member = await db.member.findFirst({
+    where: {
+      userId: session.user.id,
+      organizationId,
+      deactivated: false,
+    },
+    select: { id: true },
+  });
+
+  if (!member) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  const nowIso = new Date().toISOString();
+
   const systemPrompt = `
     You're an expert in GRC, and a helpful assistant in Comp AI,
     a platform that helps companies get compliant with frameworks
@@ -33,6 +63,12 @@ export async function POST(req: Request) {
     Keep responses concise and to the point.
 
     If you are unsure about the answer, say "I don't know" or "I don't know the answer to that question".
+
+    Important:
+    - Today's date/time is ${nowIso}.
+    - You are assisting a user inside a live application (organizationId: ${organizationId}).
+    - Prefer using available tools to fetch up-to-date org data (policies, risks, organization details) rather than guessing.
+    - If the question depends on the customer's current configuration/data and you haven't retrieved it, call the relevant tool first.
 `;
 
   const result = streamText({