JohnCosta27/Haystack
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

fix: notification system

Browse Source
This commit is contained in:
John Costa
2025-10-05 12:10:06 +01:00
parent 649cfe0b02
commit 980b42aa44
12 changed files with 191 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
backend/events.go
View File

@ -5,77 +5,22 @@ import (
"fmt"
"net/http"
"screenmark/screenmark/middleware"
"screenmark/screenmark/notifications"
"strconv"
"github.com/google/uuid"
)
const (
IMAGE_TYPE = "image"
LIST_TYPE = "list"
)
type imageNotification struct {
Type string
ImageID uuid.UUID
ImageName string
Status string
}
type listNotification struct {
Type string
ListID uuid.UUID
Name string
Status string
}
type Notification struct {
image *imageNotification
list *listNotification
}
func getImageNotification(image imageNotification) Notification {
return Notification{
image: &image,
}
}
func getListNotification(list listNotification) Notification {
return Notification{
list: &list,
}
}
func (n Notification) MarshalJSON() ([]byte, error) {
if n.image != nil {
return json.Marshal(n.image)
}
if n.list != nil {
return json.Marshal(n.list)
}
return nil, fmt.Errorf("no image or list present")
}
func (n *Notification) UnmarshalJSON(data []byte) error {
return fmt.Errorf("unimplemented")
}
/*
* TODO: We have channels open every a user sends an image.
* We never close these channels.
*
* What is a reasonable default? Close the channel after 1 minute of inactivity?
*/
func CreateEventsHandler(notifier *Notifier[Notification]) http.HandlerFunc {
func CreateEventsHandler(notifier *notifications.Notifier[notifications.Notification]) http.HandlerFunc {
counter := 0
userSplitters := make(map[string]*ChannelSplitter[Notification])
userSplitters := make(map[string]*notifications.ChannelSplitter[notifications.Notification])
return func(w http.ResponseWriter, r *http.Request) {
_userId := r.Context().Value(middleware.USER_ID).(uuid.UUID)
@ -98,7 +43,7 @@ func CreateEventsHandler(notifier *Notifier[Notification]) http.HandlerFunc {
userNotifications := notifier.Listeners[userId]
if _, exists := userSplitters[userId]; !exists {
splitter := NewChannelSplitter(userNotifications)
splitter := notifications.NewChannelSplitter(userNotifications)
userSplitters[userId] = &splitter
splitter.Listen()

6
backend/images/handler.go
View File

@ -154,7 +154,11 @@ func (h *ImageHandler) uploadImage(w http.ResponseWriter, r *http.Request) {
h.logger.Info("About to add image")
h.processor.Add(newImage)
w.WriteHeader(http.StatusOK)
// We nullify the image's data, so we're not transferring all that
// data back to the frontend.
newImage.Image = nil
middleware.WriteJsonOrError(h.logger, newImage, w)
}
func (h *ImageHandler) deleteImage(w http.ResponseWriter, r *http.Request) {

6
backend/integration_test.go
View File

@ -181,7 +181,11 @@ func setupTestContext(t *testing.T) *TestContext {
}
jwtManager := middleware.NewJwtManager([]byte("test-jwt-secret"))
router := setupRouter(db, jwtManager)
router, err := setupRouter(db, jwtManager)
if err != nil {
panic(err)
}
server := httptest.NewServer(router)
tc.db = db

5
backend/main.go
View File

@ -28,7 +28,10 @@ func main() {
panic(err)
}
router := setupRouter(db, jwtManager)
router, err := setupRouter(db, jwtManager)
if err != nil {
panic(err)
}
port, exists := os.LookupEnv("PORT")
if !exists {

38
backend/notifications/channel_splitter.go Normal file
View File

@ -0,0 +1,38 @@
package notifications
type ChannelSplitter[TNotification any] struct {
ch chan TNotification
Listeners map[string]chan TNotification
}
func (s *ChannelSplitter[TNotification]) Listen() {
go func() {
for {
select {
case msg := <-s.ch:
for _, v := range s.Listeners {
v <- msg
}
}
}
}()
}
func (s *ChannelSplitter[TNotification]) Add(id string) chan TNotification {
ch := make(chan TNotification)
s.Listeners[id] = ch
return ch
}
func (s *ChannelSplitter[TNotification]) Remove(id string) {
delete(s.Listeners, id)
}
func NewChannelSplitter[TNotification any](ch chan TNotification) ChannelSplitter[TNotification] {
return ChannelSplitter[TNotification]{
ch: ch,
Listeners: make(map[string]chan TNotification),
}
}

64
backend/notifications/image_notification.go Normal file
View File

@ -0,0 +1,64 @@
package notifications
import (
"encoding/json"
"fmt"
"github.com/google/uuid"
)
const (
IMAGE_TYPE = "image"
LIST_TYPE = "list"
)
type ImageNotification struct {
Type string
ImageID uuid.UUID
ImageName string
Status string
}
type ListNotification struct {
Type string
ListID uuid.UUID
Name string
Status string
}
type Notification struct {
image *ImageNotification
list *ListNotification
}
func GetImageNotification(image ImageNotification) Notification {
return Notification{
image: &image,
}
}
func GetListNotification(list ListNotification) Notification {
return Notification{
list: &list,
}
}
func (n Notification) MarshalJSON() ([]byte, error) {
if n.image != nil {
return json.Marshal(n.image)
}
if n.list != nil {
return json.Marshal(n.list)
}
return nil, fmt.Errorf("no image or list present")
}
func (n *Notification) UnmarshalJSON(data []byte) error {
return fmt.Errorf("unimplemented")
}

41
backend/notifications.go → backend/notifications/notifications.go
View File

@ -1,4 +1,4 @@
package main
package notifications
import (
"errors"
@ -56,42 +56,3 @@ func NewNotifier[TNotification any](bufferSize int) Notifier[TNotification] {
Listeners: make(map[string]chan TNotification),
}
}
// ----------------------------------
type ChannelSplitter[TNotification any] struct {
ch chan TNotification
Listeners map[string]chan TNotification
}
func (s *ChannelSplitter[TNotification]) Listen() {
go func() {
for {
select {
case msg := <-s.ch:
for _, v := range s.Listeners {
v <- msg
}
}
}
}()
}
func (s *ChannelSplitter[TNotification]) Add(id string) chan TNotification {
ch := make(chan TNotification)
s.Listeners[id] = ch
return ch
}
func (s *ChannelSplitter[TNotification]) Remove(id string) {
delete(s.Listeners, id)
}
func NewChannelSplitter[TNotification any](ch chan TNotification) ChannelSplitter[TNotification] {
return ChannelSplitter[TNotification]{
ch: ch,
Listeners: make(map[string]chan TNotification),
}
}

2
backend/notifications_test.go → backend/notifications/notifications_test.go
View File

@ -1,4 +1,4 @@
package main
package notifications
import (
"testing"

48
backend/processor/image.go
View File

@ -2,11 +2,13 @@ package processor
import (
"context"
"fmt"
"screenmark/screenmark/.gen/haystack/haystack/model"
"screenmark/screenmark/agents"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"screenmark/screenmark/notifications"
"sync"
"github.com/charmbracelet/log"
@ -24,6 +26,8 @@ type ImageProcessor struct {
// TODO: add the notifier here
Processor *Processor[model.Image]
notifier *notifications.Notifier[notifications.Notification]
}
func (p *ImageProcessor) setImageToProcess(ctx context.Context, image model.Image) {
@ -71,6 +75,19 @@ func (p *ImageProcessor) processImage(image model.Image) {
var wg sync.WaitGroup
wg.Add(2)
imageNotification := notifications.GetImageNotification(notifications.ImageNotification{
Type: notifications.IMAGE_TYPE,
ImageID: image.ID,
ImageName: image.ImageName,
Status: string(model.Progress_InProgress),
})
err := p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
if err != nil {
p.logger.Error("sending in progress notification", "err", err)
return
}
go func() {
p.describe(ctx, image)
wg.Done()
@ -82,9 +99,34 @@ func (p *ImageProcessor) processImage(image model.Image) {
}()
wg.Wait()
// TODO: there is some repeated code here. The ergonomicts of the notifications,
// isn't the best.
imageNotification = notifications.GetImageNotification(notifications.ImageNotification{
Type: notifications.IMAGE_TYPE,
ImageID: image.ID,
ImageName: image.ImageName,
Status: string(model.Progress_Complete),
})
err = p.notifier.SendAndCreate(image.UserID.String(), imageNotification)
if err != nil {
p.logger.Error("sending done notification", "err", err)
return
}
}
func NewImageProcessor(
logger *log.Logger,
imageModel models.ImageModel,
listModel models.StackModel,
limitsManager limits.LimitsManagerMethods,
notifier *notifications.Notifier[notifications.Notification],
) (ImageProcessor, error) {
if notifier == nil {
return ImageProcessor{}, fmt.Errorf("notifier is nil")
}
func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listModel models.StackModel, limitsManager limits.LimitsManagerMethods) ImageProcessor {
descriptionAgent := agents.NewDescriptionAgent(logger, imageModel)
stackAgent := agents.NewListAgent(logger, listModel, limitsManager)
@ -93,9 +135,11 @@ func NewImageProcessor(logger *log.Logger, imageModel models.ImageModel, listMod
logger: logger,
descriptionAgent: descriptionAgent,
stackAgent: stackAgent,
notifier: notifier,
}
imageProcessor.Processor = NewProcessor(int(IMAGE_PROCESS_AT_A_TIME), imageProcessor.processImage)
return imageProcessor
return imageProcessor, nil
}

16
backend/router.go
View File

@ -2,12 +2,14 @@ package main
import (
"database/sql"
"fmt"
"os"
"screenmark/screenmark/agents/client"
"screenmark/screenmark/auth"
"screenmark/screenmark/images"
"screenmark/screenmark/limits"
"screenmark/screenmark/models"
"screenmark/screenmark/notifications"
"screenmark/screenmark/processor"
"screenmark/screenmark/stacks"
@ -25,22 +27,26 @@ func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (cli
return client.ImageInfo, nil
}
func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) chi.Router {
func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) (chi.Router, error) {
limitsManager := limits.CreateLimitsManager(db)
imageModel := models.NewImageModel(db)
stackModel := models.NewStackModel(db)
notifier := notifications.NewNotifier[notifications.Notification](10)
imageProcessorLogger := createLogger("Image Processor", os.Stdout)
imageProcessor := processor.NewImageProcessor(imageProcessorLogger, imageModel, stackModel, limitsManager)
imageProcessor, err := processor.NewImageProcessor(imageProcessorLogger, imageModel, stackModel, limitsManager, &notifier)
if err != nil {
return nil, fmt.Errorf("processor: %w", err)
}
go imageProcessor.Processor.Work()
stackHandler := stacks.CreateStackHandler(db, limitsManager, jwtManager)
authHandler := auth.CreateAuthHandler(db, jwtManager)
imageHandler := images.CreateImageHandler(db, limitsManager, jwtManager, imageProcessor.Processor)
notifier := NewNotifier[Notification](10)
r := chi.NewRouter()
r.Use(middleware.Logger)
@ -56,5 +62,5 @@ func setupRouter(db *sql.DB, jwtManager *ourmiddleware.JwtManager) chi.Router {
r.Get("/", CreateEventsHandler(&notifier))
})
return r
return r, nil
}

4
frontend/src/contexts/Notifications.tsx
View File

@ -39,6 +39,8 @@ export const Notifications = (onCompleteImage: () => void) => {
const [accessToken] = createResource(getAccessToken);
const dataEventListener = (e: MessageEvent<unknown>) => {
debugger;
if (typeof e.data !== "string") {
console.error("Error type is not string");
return;
@ -98,7 +100,7 @@ export const Notifications = (onCompleteImage: () => void) => {
upsertImageProcessing(
Object.fromEntries(
images.filter(i => i.Status !== 'complete').map((i) => [
images.filter(i => i.Status === 'complete').map((i) => [
i.ID,
{
Type: "image",

22
frontend/src/network/index.ts
View File

@ -78,17 +78,10 @@ const getBaseAuthorizedRequest = async ({
method,
});
};
const sendImageResponseValidator = strictObject({
ID: pipe(string(), uuid()),
ImageID: pipe(string(), uuid()),
UserID: pipe(string(), uuid()),
Status: string(),
});
export const sendImageFile = async (
imageName: string,
file: File,
): Promise<InferOutput<typeof sendImageResponseValidator>> => {
): Promise<InferOutput<typeof imageValidator>> => {
const request = await getBaseAuthorizedRequest({
path: `images/${imageName}`,
body: file,
@ -98,7 +91,7 @@ export const sendImageFile = async (
request.headers.set("Content-Type", "application/oclet-stream");
const res = await fetch(request).then((res) => res.json());
const parsedRes = safeParse(sendImageResponseValidator, res);
const parsedRes = safeParse(imageValidator, res);
if (!parsedRes.success) {
console.log(parsedRes.issues)
@ -146,7 +139,7 @@ export class ImageLimitReached extends Error {
export const sendImage = async (
imageName: string,
base64Image: string,
): Promise<InferOutput<typeof sendImageResponseValidator>> => {
): Promise<InferOutput<typeof imageValidator>> => {
const request = await getBaseAuthorizedRequest({
path: `images/${imageName}`,
body: base64Image,
@ -162,16 +155,16 @@ export const sendImage = async (
const res = await rawRes.json();
const parsedRes = safeParse(sendImageResponseValidator, res);
const parsedRes = safeParse(imageValidator, res);
if (!parsedRes.success) {
console.log(parsedRes.issues)
console.log("Parsing issues: ", parsedRes.issues)
throw new Error(JSON.stringify(parsedRes.issues));
}
return parsedRes.output;
};
const userImageValidator = strictObject({
const imageValidator = strictObject({
ID: pipe(string(), uuid()),
CreatedAt: string(),
UserID: pipe(string(), uuid()),
@ -181,7 +174,10 @@ const userImageValidator = strictObject({
ImageName: string(),
Status: union([literal('not-started'), literal('in-progress'), literal('complete')]),
})
const userImageValidator = strictObject({
...imageValidator.entries,
ImageStacks: pipe(nullable(array(
strictObject({
ID: pipe(string(), uuid()),