feat: ComfyUI fallback + session fixes

.gitignore

@@ -126,3 +126,4 @@ apps/stage-pocket/ios/buildServer.json
 .ghfs/
 .docs/superpowers
 apps/stage-tamagotchi/electron.vite.config.*.mjs
+.env

apps/server/docker-compose.yml

@@ -70,6 +70,23 @@ services:
         required: false
     restart: unless-stopped
 
+  nsfw-image-consumer:
+    build:
+      context: ../..
+      dockerfile: apps/server/Dockerfile
+    command: ['pnpm', '-F', '@proj-airi/server', 'run', 'server', 'nsfw-image-consumer']
+    depends_on:
+      db:
+        condition: service_healthy
+      redis:
+        condition: service_healthy
+    env_file:
+      - path: .env
+        required: false
+      - path: .env.local
+        required: false
+    restart: unless-stopped
+
 volumes:
   db_data:
     driver: local

apps/server/drizzle/0008_lazy_wallflower.sql

@@ -0,0 +1,14 @@
+ALTER TABLE "characters"
+ADD COLUMN "visibility" text DEFAULT 'private' NOT NULL;
+--> statement-breakpoint
+ALTER TABLE "characters"
+ADD COLUMN "nsfw_enabled" boolean DEFAULT false NOT NULL;
+--> statement-breakpoint
+ALTER TABLE "characters"
+ADD COLUMN "nsfw_level" text DEFAULT 'none' NOT NULL;
+--> statement-breakpoint
+ALTER TABLE "characters"
+ADD COLUMN "relationship_mode" text DEFAULT 'companion' NOT NULL;
+--> statement-breakpoint
+ALTER TABLE "characters"
+ADD COLUMN "persona_profile" jsonb DEFAULT '{}'::jsonb NOT NULL;

apps/server/drizzle/0009_nsfw_user_gating.sql

@@ -0,0 +1,8 @@
+ALTER TABLE "user"
+  ADD COLUMN "adult_verified" boolean DEFAULT false NOT NULL;
+
+ALTER TABLE "user"
+  ADD COLUMN "allow_sensitive_content" boolean DEFAULT false NOT NULL;
+
+ALTER TABLE "user"
+  ADD COLUMN "content_tier" text DEFAULT 'standard' NOT NULL;

apps/server/drizzle/0010_nsfw_media.sql

@@ -0,0 +1,36 @@
+CREATE TABLE "image_jobs" (
+  "id" text PRIMARY KEY NOT NULL,
+  "user_id" text NOT NULL,
+  "character_id" text NOT NULL,
+  "route" text NOT NULL,
+  "status" text DEFAULT 'queued' NOT NULL,
+  "prompt" text NOT NULL,
+  "negative_prompt" text NOT NULL,
+  "scene_type" text,
+  "tags" text[] DEFAULT '{}' NOT NULL,
+  "params" jsonb DEFAULT '{}'::jsonb NOT NULL,
+  "result_media_id" text,
+  "error_message" text,
+  "created_at" timestamp DEFAULT now() NOT NULL,
+  "updated_at" timestamp DEFAULT now() NOT NULL
+);
+
+CREATE TABLE "gallery_items" (
+  "id" text PRIMARY KEY NOT NULL,
+  "user_id" text NOT NULL,
+  "character_id" text NOT NULL,
+  "image_job_id" text,
+  "media_id" text,
+  "title" text,
+  "prompt" text NOT NULL,
+  "negative_prompt" text NOT NULL,
+  "scene_type" text,
+  "tags" text[] DEFAULT '{}' NOT NULL,
+  "created_at" timestamp DEFAULT now() NOT NULL,
+  "updated_at" timestamp DEFAULT now() NOT NULL
+);
+
+ALTER TABLE "gallery_items"
+  ADD CONSTRAINT "gallery_items_image_job_id_image_jobs_id_fk"
+  FOREIGN KEY ("image_job_id") REFERENCES "public"."image_jobs"("id")
+  ON DELETE set null ON UPDATE no action;

apps/server/drizzle/meta/_journal.json

@@ -57,6 +57,27 @@
       "when": 1774632446757,
       "tag": "0007_red_nicolaos",
       "breakpoints": true
+    },
+    {
+      "idx": 8,
+      "version": "7",
+      "when": 1774800000000,
+      "tag": "0008_lazy_wallflower",
+      "breakpoints": true
+    },
+    {
+      "idx": 9,
+      "version": "7",
+      "when": 1774803600000,
+      "tag": "0009_nsfw_user_gating",
+      "breakpoints": true
+    },
+    {
+      "idx": 10,
+      "version": "7",
+      "when": 1774807200000,
+      "tag": "0010_nsfw_media",
+      "breakpoints": true
     }
   ]
-}
+}

apps/server/package.json

@@ -9,6 +9,7 @@
     "dev": "pnpm run apply:env -- tsx --watch src/bin/run.ts api",
     "start": "pnpm run apply:env -- tsx src/bin/run.ts api",
     "server": "pnpm run apply:env -- tsx src/bin/run.ts",
+    "nsfw:image-consumer": "pnpm run apply:env -- tsx src/bin/run.ts nsfw-image-consumer",
     "build": "tsc -b",
     "typecheck": "tsc --noEmit",
     "db:generate": "drizzle-kit generate",

apps/server/src/app.ts

@@ -10,7 +10,10 @@ import type { ChatService } from './services/chats'
 import type { ConfigKVService } from './services/config-kv'
 import type { FluxService } from './services/flux'
 import type { FluxTransactionService } from './services/flux-transaction'
+import type { Database } from './libs/db'
 import type { ProviderService } from './services/providers'
+import type { NsfwMediaService } from './services/nsfw-media'
+import type { NsfwImageEvent } from './services/nsfw-image-events'
 import type { StripeService } from './services/stripe'
 import type { HonoEnv } from './types/hono'
 
@@ -40,8 +43,11 @@ import { createChatRoutes } from './routes/chats'
 import { createFluxRoutes } from './routes/flux'
 import { createV1CompletionsRoutes } from './routes/openai/v1'
 import { createProviderRoutes } from './routes/providers'
+import { createNsfwMediaRoutes } from './routes/nsfw-media'
 import { createStripeRoutes } from './routes/stripe'
+import { createUserRoutes } from './routes/users'
 import { createBillingMq } from './services/billing/billing-events'
+import { createNsfwImageMq } from './services/nsfw-image-events'
 import { createBillingService } from './services/billing/billing-service'
 import { createCharacterService } from './services/characters'
 import { createChatService } from './services/chats'
@@ -50,20 +56,24 @@ import { createFluxService } from './services/flux'
 import { createFluxTransactionService } from './services/flux-transaction'
 import { createProviderService } from './services/providers'
 import { createRequestLogService } from './services/request-log'
+import { createNsfwMediaService } from './services/nsfw-media'
 import { createStripeService } from './services/stripe'
 import { ApiError, createInternalError, createUnauthorizedError } from './utils/error'
 import { getTrustedOrigin } from './utils/origin'
 
 interface AppDeps {
+  db: Database
   auth: ReturnType<typeof createAuth>
   characterService: CharacterService
   chatService: ChatService
   providerService: ProviderService
+  nsfwMediaService: NsfwMediaService
   fluxService: FluxService
   fluxTransactionService: FluxTransactionService
   stripeService: StripeService
   billingService: BillingService
   billingMq: MqService<BillingEvent>
+  nsfwImageMq: MqService<NsfwImageEvent>
   configKV: ConfigKVService
   redis: Redis
   env: Env
@@ -191,6 +201,16 @@ async function buildApp(deps: AppDeps) {
      */
     .route('/api/v1/providers', createProviderRoutes(deps.providerService))
 
+    /**
+     * Current user settings routes.
+     */
+    .route('/api/v1/users', createUserRoutes(deps.db))
+
+    /**
+     * NSFW media routes.
+     */
+    .route('/api/v1/nsfw', createNsfwMediaRoutes(deps.nsfwMediaService, deps.nsfwImageMq))
+
     /**
      * Chat routes are handled by the chat service.
      */
@@ -321,6 +341,18 @@ export async function createApp() {
     build: ({ dependsOn }) => createProviderService(dependsOn.db),
   })
 
+  const nsfwMediaService = injeca.provide('services:nsfw-media', {
+    dependsOn: { db },
+    build: ({ dependsOn }) => createNsfwMediaService(dependsOn.db),
+  })
+
+  const nsfwImageMq = injeca.provide('services:nsfwImageMq', {
+    dependsOn: { redis, env: parsedEnv },
+    build: ({ dependsOn }) => createNsfwImageMq(dependsOn.redis, {
+      stream: dependsOn.env.NSFW_IMAGE_EVENTS_STREAM,
+    }),
+  })
+
   const chatService = injeca.provide('services:chats', {
     dependsOn: { db, otel },
     build: ({ dependsOn }) => createChatService(dependsOn.db, dependsOn.otel?.engagement),
@@ -358,27 +390,32 @@ export async function createApp() {
     characterService,
     chatService,
     providerService,
+    nsfwMediaService,
     fluxService,
     fluxTransactionService,
     requestLogService,
     stripeService,
     billingService,
     billingMq,
+    nsfwImageMq,
     configKV,
     redis,
     env: parsedEnv,
     otel,
   })
   const { app, injectWebSocket } = await buildApp({
+    db: resolved.db,
     auth: resolved.auth,
     characterService: resolved.characterService,
     chatService: resolved.chatService,
-    providerService: resolved.providerService,
-    fluxService: resolved.fluxService,
+      providerService: resolved.providerService,
+      nsfwMediaService: resolved.nsfwMediaService,
+      fluxService: resolved.fluxService,
     fluxTransactionService: resolved.fluxTransactionService,
     stripeService: resolved.stripeService,
     billingService: resolved.billingService,
     billingMq: resolved.billingMq,
+    nsfwImageMq: resolved.nsfwImageMq,
     configKV: resolved.configKV,
     redis: resolved.redis,
     env: resolved.env,

apps/server/src/bin/run-nsfw-image-consumer.ts

@@ -0,0 +1,94 @@
+import process, { pid } from 'node:process'
+
+import { initLogger, LoggerFormat, LoggerLevel, useLogger } from '@guiiai/logg'
+
+import { createDrizzle, migrateDatabase } from '../libs/db'
+import { parseEnv } from '../libs/env'
+import { initializeExternalDependency } from '../libs/external-dependency'
+import { createMqWorker } from '../libs/mq'
+import { createRedis } from '../libs/redis'
+import { createNsfwImageConsumerHandler } from '../services/nsfw-image-consumer-handler'
+import { createNsfwImageMq } from '../services/nsfw-image-events'
+import { createNsfwMediaService } from '../services/nsfw-media'
+
+export async function runNsfwImageConsumer(): Promise<void> {
+  initLogger(LoggerLevel.Debug, LoggerFormat.Pretty)
+
+  const env = parseEnv(process.env)
+  const logger = useLogger('nsfw-image-consumer').useGlobalConfig()
+  const { db, pool } = await initializeExternalDependency(
+    'Database',
+    logger,
+    async (attempt) => {
+      const connection = createDrizzle(env)
+
+      try {
+        await connection.db.execute('SELECT 1')
+        logger.log(`Connected to database on attempt ${attempt}`)
+        await migrateDatabase(connection.db)
+        logger.log(`Applied schema on attempt ${attempt}`)
+        return connection
+      }
+      catch (error) {
+        await connection.pool.end()
+        throw error
+      }
+    },
+  )
+  const redis = await initializeExternalDependency(
+    'Redis',
+    logger,
+    async (attempt) => {
+      const instance = createRedis(env.REDIS_URL)
+
+      try {
+        await instance.connect()
+        logger.log(`Connected to Redis on attempt ${attempt}`)
+        return instance
+      }
+      catch (error) {
+        instance.disconnect()
+        throw error
+      }
+    },
+  )
+
+  const abortController = new AbortController()
+  const consumer = env.NSFW_IMAGE_EVENTS_CONSUMER_NAME ?? `nsfw-image-consumer-${pid}`
+
+  const shutdown = (signalName: string) => {
+    if (abortController.signal.aborted) {
+      return
+    }
+
+    logger.withFields({ signalName }).log('Stopping NSFW image consumer')
+    abortController.abort()
+  }
+
+  process.once('SIGINT', () => shutdown('SIGINT'))
+  process.once('SIGTERM', () => shutdown('SIGTERM'))
+
+  try {
+    const mq = createNsfwImageMq(redis, {
+      stream: env.NSFW_IMAGE_EVENTS_STREAM,
+    })
+
+    const mediaService = createNsfwMediaService(db)
+    const handler = createNsfwImageConsumerHandler(mediaService, env)
+    const worker = createMqWorker(mq)
+
+    await worker.run({
+      group: 'nsfw-image-consumer',
+      consumer,
+      signal: abortController.signal,
+      batchSize: env.NSFW_IMAGE_EVENTS_BATCH_SIZE,
+      blockMs: env.NSFW_IMAGE_EVENTS_BLOCK_MS,
+      minIdleTimeMs: env.NSFW_IMAGE_EVENTS_MIN_IDLE_MS,
+      onMessage: message => handler.handleMessage(message.event),
+    })
+  }
+  finally {
+    await redis.quit()
+    await pool.end()
+  }
+}

apps/server/src/bin/run.ts

@@ -9,8 +9,9 @@ import { cac } from 'cac'
 import { runApiServer } from '../app'
 import { errorMessageFromUnknown } from '../utils/error-message'
 import { runBillingConsumer } from './run-billing-consumer'
+import { runNsfwImageConsumer } from './run-nsfw-image-consumer'
 
-const serverRoles = ['api', 'billing-consumer'] as const
+const serverRoles = ['api', 'billing-consumer', 'nsfw-image-consumer'] as const
 
 type ServerRole = typeof serverRoles[number]
 
@@ -21,6 +22,9 @@ async function runServerRole(role: ServerRole): Promise<void> {
       return
     case 'billing-consumer':
       await runBillingConsumer()
+      return
+    case 'nsfw-image-consumer':
+      await runNsfwImageConsumer()
   }
 }
 
@@ -36,6 +40,10 @@ export function createServerCli() {
     .command('billing-consumer', 'Start the billing events consumer (transactions, audit, request logs)')
     .action(() => runServerRole('billing-consumer'))
 
+  cli
+    .command('nsfw-image-consumer', 'Start the NSFW image jobs consumer (ComfyUI execution and reconciliation)')
+    .action(() => runServerRole('nsfw-image-consumer'))
+
   cli.help()
 
   return cli