diff --git a/apps/server/src/core/ai-chat/ai-chat.controller.ts b/apps/server/src/core/ai-chat/ai-chat.controller.ts
index 6abaae72..04a255c0 100644
--- a/apps/server/src/core/ai-chat/ai-chat.controller.ts
+++ b/apps/server/src/core/ai-chat/ai-chat.controller.ts
@@ -21,7 +21,13 @@ import { JwtAuthGuard } from '../../common/guards/jwt-auth.guard';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
 import { AuthWorkspace } from '../../common/decorators/auth-workspace.decorator';
 import { SkipTransform } from '../../common/decorators/skip-transform.decorator';
-import { AiChat, User, Workspace } from '@docmost/db/types/entity.types';
+import {
+  AiChat,
+  AiChatMessage,
+  AiChatRun,
+  User,
+  Workspace,
+} from '@docmost/db/types/entity.types';
 import { PaginationOptions } from '@docmost/db/pagination/pagination-options';
 import { AiChatRepo } from '@docmost/db/repos/ai-chat/ai-chat.repo';
 import { AiChatMessageRepo } from '@docmost/db/repos/ai-chat/ai-chat-message.repo';
@@ -160,7 +166,7 @@ export class AiChatController {
     @Body() dto: GetRunDto,
     @AuthUser() user: User,
     @AuthWorkspace() workspace: Workspace,
-  ): Promise<{ run: unknown; message: unknown }> {
+  ): Promise<{ run: AiChatRun | null; message: AiChatMessage | null }> {
     await this.assertOwnedChat(dto.chatId, user, workspace);
     const run = await this.aiChatRunService.getLatestForChat(
       dto.chatId,