frame image

Author: mariusandra

backend/app/api/frames.py

@@ -1,8 +1,10 @@
+from datetime import datetime, timedelta
 import io
 import json
 import os
 import shlex
 import asyncio
+from jose import JWTError, jwt
 from http import HTTPStatus
 from tempfile import NamedTemporaryFile
 
@@ -19,7 +21,9 @@
 from app.codegen.scene_nim import write_scene_nim
 from app.utils.ssh_utils import get_ssh_connection, exec_command, remove_ssh_connection
 from scp import SCPClient
-from . import private_api
+
+from app.api.auth import ALGORITHM, SECRET_KEY, get_current_user
+from . import private_api, public_api
 
 
 @private_api.get("/frames")
@@ -58,8 +62,30 @@ async def api_frame_get_logs(id: int, db: Session = Depends(get_db)):
                             status_code=HTTPStatus.NOT_FOUND)
 
 
-@private_api.get("/frames/{id}/image")
-async def api_frame_get_image(id: int, request: Request, db: Session = Depends(get_db)):
+@private_api.get("/frames/{id}/image_link")
+async def get_image_link(id: int, user=Depends(get_current_user)):
+    expire_minutes = 5
+    now = datetime.utcnow()
+    expire = now + timedelta(minutes=expire_minutes)
+    to_encode = {"sub": str(id), "exp": expire}
+    token = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+
+    expires_in = int((expire - now).total_seconds())
+
+    return {
+        "url": f"/api/frames/{id}/image?token={token}",
+        "expires_in": expires_in
+    }
+
+@public_api.get("/frames/{id}/image")
+async def api_frame_get_image(id: int, token: str, request: Request, db: Session = Depends(get_db)):
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        if payload.get("sub") != str(id):
+            return JSONResponse({"error": "Unauthorized"}, status_code=401)
+    except JWTError:
+        return JSONResponse({"error": "Unauthorized"}, status_code=401)
+
     frame = db.query(Frame).get(id)
     if frame is None:
         return JSONResponse(content={'error': 'Frame not found'}, status_code=HTTPStatus.NOT_FOUND)
@@ -70,6 +96,8 @@ async def api_frame_get_image(id: int, request: Request, db: Session = Depends(g
         url += "?k=" + frame.frame_access_key
 
     try:
+        # We can check if a cached version exists
+        # TODO: no request object here
         if request.query_params.get('t') == '-1':
             last_image = await redis.get(cache_key)
             if last_image:
@@ -82,9 +110,6 @@ async def api_frame_get_image(id: int, request: Request, db: Session = Depends(g
             await redis.set(cache_key, response.content, ex=86400 * 30)
             return Response(content=response.content, media_type='image/png')
         else:
-            last_image = await redis.get(cache_key)
-            if last_image:
-                return Response(content=last_image, media_type='image/png')
             return JSONResponse(content={"error": "Unable to fetch image"}, status_code=response.status_code)
 
     except httpx.ReadTimeout:
@@ -94,7 +119,6 @@ async def api_frame_get_image(id: int, request: Request, db: Session = Depends(g
         return JSONResponse(content={'error': 'Internal Server Error', 'message': str(e)},
                             status_code=HTTPStatus.INTERNAL_SERVER_ERROR)
 
-
 @private_api.get("/frames/{id}/state")
 async def api_frame_get_state(id: int, db: Session = Depends(get_db)):
     frame = db.query(Frame).get(id)

frontend/src/components/FrameImage.tsx

@@ -0,0 +1,90 @@
+import { useActions, useValues } from 'kea'
+import clsx from 'clsx'
+import { framesModel } from '../models/framesModel'
+import { useEffect, useState } from 'react'
+
+export interface FrameImageProps extends React.HTMLAttributes<HTMLDivElement> {
+  frameId: number
+  className?: string
+  /** If true, user can click on the image to request a refresh of the signed URL */
+  refreshable?: boolean
+}
+
+/**
+ * Consolidated Image component:
+ * - Takes a frameId
+ * - Uses framesModel to get and update the frame's image
+ * - Shows loading states based on image load or frame readiness
+ * - Optionally allows clicking the image container to refresh the image link if `refreshable` is true
+ */
+export function FrameImage({ frameId, className, refreshable = true, ...props }: FrameImageProps) {
+  const { getFrameImage, frames } = useValues(framesModel)
+  const { updateFrameImage } = useActions(framesModel)
+
+  const [isLoading, setIsLoading] = useState(true)
+
+  const imageUrl = getFrameImage(frameId)
+  const frame = frames[frameId]
+
+  useEffect(() => {
+    updateFrameImage(frameId, false)
+  }, [!!imageUrl])
+
+  useEffect(() => {
+    // Whenever the image URL changes, we consider the image as loading again
+    // because the <img> will re-attempt to load the new URL.
+    if (imageUrl) {
+      setIsLoading(true)
+    }
+  }, [imageUrl])
+
+  console.log({ imageUrl })
+
+  // Determine if we should show the fade-in-out or loading cursor
+  const visiblyLoading = (isLoading || frame?.status !== 'ready') && frame?.interval > 5
+
+  const handleRefreshClick = () => {
+    if (refreshable) {
+      updateFrameImage(frameId)
+    }
+  }
+
+  return (
+    <div
+      className={clsx(
+        'max-w-full max-h-full w-full h-full flex items-center justify-center',
+        visiblyLoading ? 'continuous-fade-in-out' : null,
+        visiblyLoading ? 'cursor-wait' : refreshable ? 'cursor-pointer' : 'cursor-default',
+        className
+      )}
+      onClick={handleRefreshClick}
+      title={refreshable ? 'Click to refresh' : undefined}
+      {...props}
+    >
+      {frame && (
+        <img
+          className={clsx('rounded-lg', refreshable ? 'rounded-tl-none max-w-full max-h-full' : 'w-full')}
+          src={imageUrl ?? undefined}
+          onLoad={() => setIsLoading(false)}
+          onError={() => setIsLoading(false)}
+          style={{
+            ...(frame.width && frame.height
+              ? {
+                  aspectRatio:
+                    frame.rotate === 90 || frame.rotate === 270
+                      ? `${frame.height} / ${frame.width}`
+                      : `${frame.width} / ${frame.height}`,
+                  maxWidth: '100%',
+                  maxHeight: '100%',
+                  // If you need a fixed width/height based on rotation:
+                  // width: frame.rotate === 90 || frame.rotate === 270 ? frame.height : frame.width,
+                  // height: 'auto',
+                }
+              : {}),
+          }}
+          alt=""
+        />
+      )}
+    </div>
+  )
+}

frontend/src/models/framesModel.tsx

@@ -1,14 +1,17 @@
 import { actions, afterMount, connect, kea, listeners, path, reducers, selectors } from 'kea'
-
 import { loaders } from 'kea-loaders'
 import { FrameScene, FrameType } from '../types'
 import { socketLogic } from '../scenes/socketLogic'
-
 import type { framesModelType } from './framesModelType'
 import { router } from 'kea-router'
 import { sanitizeScene } from '../scenes/frame/frameLogic'
 import { apiFetch } from '../utils/apiFetch'
 
+export interface FrameImageInfo {
+  url: string
+  expiresAt: number
+}
+
 export const framesModel = kea<framesModelType>([
   connect({ logic: [socketLogic] }),
   path(['src', 'models', 'framesModel']),
@@ -18,8 +21,10 @@ export const framesModel = kea<framesModelType>([
     redeployFrame: (id: number) => ({ id }),
     restartFrame: (id: number) => ({ id }),
     renderFrame: (id: number) => ({ id }),
-    updateFrameImage: (id: number) => ({ id }),
+    updateFrameImage: (id: number, force = true) => ({ id, force }),
     deleteFrame: (id: number) => ({ id }),
+    setFrameImageInfo: (id: number, imageInfo: FrameImageInfo) => ({ id, imageInfo }),
+    updateFrameImageTimestamp: (id: number) => ({ id }),
   }),
   loaders(({ values }) => ({
     frames: [
@@ -29,13 +34,16 @@ export const framesModel = kea<framesModelType>([
           try {
             const response = await apiFetch(`/api/frames/${id}`)
             if (!response.ok) {
-              throw new Error('Failed to fetch logs')
+              throw new Error('Failed to fetch frame')
             }
             const data = await response.json()
             const frame = data.frame as FrameType
             return {
               ...values.frames,
-              frame: { ...frame, scenes: frame.scenes?.map((scene) => sanitizeScene(scene as FrameScene, frame)) },
+              [frame.id]: {
+                ...frame,
+                scenes: frame.scenes?.map((scene) => sanitizeScene(scene as FrameScene, frame)),
+              },
             }
           } catch (error) {
             console.error(error)
@@ -49,7 +57,16 @@ export const framesModel = kea<framesModelType>([
               throw new Error('Failed to fetch frames')
             }
             const data = await response.json()
-            return Object.fromEntries((data.frames as FrameType[]).map((frame) => [frame.id, frame]))
+            const framesDict = Object.fromEntries(
+              (data.frames as FrameType[]).map((frame) => [
+                frame.id,
+                {
+                  ...frame,
+                  scenes: frame.scenes?.map((scene) => sanitizeScene(scene as FrameScene, frame)),
+                },
+              ])
+            )
+            return framesDict
           } catch (error) {
             console.error(error)
             return values.frames
@@ -60,7 +77,7 @@ export const framesModel = kea<framesModelType>([
   })),
   reducers(() => ({
     frames: [
-      {} as Record<string, FrameType>,
+      {} as Record<number, FrameType>,
       {
         [socketLogic.actionTypes.newFrame]: (state, { frame }) => ({ ...state, [frame.id]: frame }),
         [socketLogic.actionTypes.updateFrame]: (state, { frame }) => ({ ...state, [frame.id]: frame }),
@@ -71,11 +88,20 @@ export const framesModel = kea<framesModelType>([
         },
       },
     ],
+    frameImageInfos: [
+      {} as Record<number, FrameImageInfo>,
+      {
+        setFrameImageInfo: (state, { id, imageInfo }) => ({ ...state, [id]: imageInfo }),
+      },
+    ],
     frameImageTimestamps: [
-      {} as Record<string, number>,
+      {} as Record<number, number>,
       {
-        updateFrameImage: (state, { id }) =>
-          state[id] === Math.floor(Date.now() / 1000) ? state : { ...state, [id]: Math.floor(Date.now() / 1000) },
+        updateFrameImageTimestamp: (state, { id }) => {
+          const nowSeconds = Math.floor(Date.now() / 1000)
+          // Only update if it's different, to ensure a re-render
+          return state[id] === nowSeconds ? state : { ...state, [id]: nowSeconds }
+        },
       },
     ],
   })),
@@ -88,18 +114,31 @@ export const framesModel = kea<framesModelType>([
         ) as FrameType[],
     ],
     getFrameImage: [
-      (s) => [s.frameImageTimestamps],
-      (frameImageTimestamps) => {
-        return (id) => {
-          return `/api/frames/${id}/image?t=${frameImageTimestamps[id] ?? -1}`
+      (s) => [s.frameImageInfos, s.frameImageTimestamps],
+      (frameImageInfos, frameImageTimestamps) => {
+        return (id: number) => {
+          const info = frameImageInfos[id]
+          if (!info) {
+            return null
+          }
+
+          const now = Math.floor(Date.now() / 1000)
+          if (!info.expiresAt || now >= info.expiresAt) {
+            // URL expired or invalid
+            return null
+          }
+
+          const timestamp = frameImageTimestamps[id] ?? -1
+          // Append timestamp as a cache-buster
+          return `${info.url}${info.url.includes('?') ? '&' : '?'}t=${timestamp}`
         }
       },
     ],
   }),
   afterMount(({ actions }) => {
     actions.loadFrames()
   }),
-  listeners(({ props, actions }) => ({
+  listeners(({ actions, values }) => ({
     renderFrame: async ({ id }) => {
       await apiFetch(`/api/frames/${id}/event/render`, { method: 'POST' })
     },
@@ -115,6 +154,33 @@ export const framesModel = kea<framesModelType>([
         router.actions.push('/')
       }
     },
+    updateFrameImage: async ({ id, force }) => {
+      // Check if we have a valid URL
+      const imageUrl = values.getFrameImage(id)
+      if (imageUrl) {
+        // The URL is still valid, no need to refetch new signed URL
+        // Just update timestamp to refresh (force reload)
+        if (force) {
+          actions.updateFrameImageTimestamp(id)
+        }
+        return
+      }
+
+      // Need a new signed URL
+      const resp = await apiFetch(`/api/frames/${id}/image_link`)
+      if (resp.ok) {
+        const data = await resp.json()
+        const expiresAt = Math.floor(Date.now() / 1000) + data.expires_in
+        const imageInfo: FrameImageInfo = { url: data.url, expiresAt }
+        actions.setFrameImageInfo(id, imageInfo)
+        // Update timestamp to ensure a new request even if the URL is same
+        if (force) {
+          actions.updateFrameImageTimestamp(id)
+        }
+      } else {
+        console.error('Failed to get image link for frame', id)
+      }
+    },
     [socketLogic.actionTypes.newLog]: ({ log }) => {
       if (log.type === 'webhook') {
         const parsed = JSON.parse(log.line)