(Production only) AIServer randomly shutting down (AI chatter stops responding mid-stream or between messages)

I've been having a tough time figuring this out. On coderealtime.com I am seeing an issue where the AI chatbot connection appears to stop responding after a period of time.

Curious if anyone is seeing the same thing or has an idea of how to debug / keep the connection of the AI chatbot to the room alive.

Seeing this happen whether I enable hibernation or not
It seems to last for like 5 minutes / 3-5 inferences
Other chats are not impacted
I see nothing in the logs at the time it seems to stop responding

My ai.ts (prompts etc stripped out) below:

import type * as Party from "partykit/server";
import { nanoid } from "nanoid";
import type { Message, ChatMessage, UserMessage } from "./utils/message";
import type { User } from "./utils/auth";
import { getChatCompletionResponse, AIMessage } from "./utils/openai";
import { notFound } from "next/navigation";
import { error, ok } from "./utils/response";

const PROMPT = `// Placeholder for prompt template`;

export const AI_USERNAME = "AI Bot";
export const AI_USER: User = {
  username: AI_USERNAME,
  expires: new Date(2099, 0, 1).toISOString(),
};

export default class AIServer implements Party.Server {
  constructor(public party: Party.Party) {
    console.log("Constructing new server.");
    setInterval(() => {
      console.log("Sending heartbeat");
      this.party.broadcast("heartbeat");
    }, 15 * 1000);
  }

  async onRequest(req: Party.Request) {
    if (req.method !== "POST") return notFound();

    const { roomId, botId, action } = await req.json<{
      roomId: string;
      botId: string;
      action: string;
    }>();
    if (action !== "connect") return notFound();

    const chatRoom = this.party.context.parties.chatroom.get(roomId);
    const socket = await chatRoom.socket("/?_pk=" + botId);

    this.simulateUser(chatRoom, botId, socket);
    return ok();
  }

  async simulateUser(chatRoom: Party.Stub, botId: string, socket: any) {
    let messages: Message[] = [];
    let activeInferences: string[] = [];

    socket.addEventListener("message", async (message) => {
      const data = JSON.parse(message.data as string) as ChatMessage;

      if (data.type === "sync") {
        messages = data.messages;
      } else if (data.type === "edit") {
        messages = messages.map((m) => (m.id === data.id ? data : m));
      } else if (data.type === "new") {
        messages.push(data);
        const fromHuman = data.from.id !== AI_USERNAME && data.from.id !== "system";

        if (fromHuman) {
          let prompt: AIMessage[] = [
            { role: "system", content: PROMPT },
            { role: "user", content: messages.map((message) => `${message.from.id}: ${message.text}`).join("\n") },
          ];

          const id = nanoid();
          let text = "";

          console.log("Getting completion response for", id);
          activeInferences.push(id);
          await getChatCompletionResponse(
            this.party.env,
            prompt,
            () => {
              socket.send(JSON.stringify(<UserMessage>{ type: "new", id, text }));
            },
            (token) => {
              text += token;
              socket.send(JSON.stringify(<UserMessage>{ type: "edit", id, text }));
            }
          );
          activeInferences = activeInferences.filter((inference) => inference !== id);
          if (activeInferences.length === 0) {
            console.log("Finished completion response for", id);
            console.log("No other active inferences, would have rebooted server");
          }
        }
      }
    });

    setInterval(() => {
      if (socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({ type: "ping" }));
      }
    }, 30000);
  }
}

partykit / partykit-nextjs-chat-template

(Production only) AIServer randomly shutting down (AI chatter stops responding mid-stream or between messages) #13