Skip to content

fix: wire Firestore client and memory store into main.go (#126) #129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ComBba merged 2 commits into from
Feb 25, 2026
Merged

fix: wire Firestore client and memory store into main.go (#126) #129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bcfa938
fix: wire Firestore client and memory store into main.go
ComBba Feb 25, 2026
38c8625
fix: address review feedback for Firestore wiring
ComBba Feb 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion cmd/server/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,15 @@ import (
"syscall"
"time"

"cloud.google.com/go/firestore"
"cloud.google.com/go/storage"
"github.com/Two-Weeks-Team/missless/internal/auth"
"github.com/Two-Weeks-Team/missless/internal/config"
"github.com/Two-Weeks-Team/missless/internal/handler"
"github.com/Two-Weeks-Team/missless/internal/media"
"github.com/Two-Weeks-Team/missless/internal/memory"
"github.com/Two-Weeks-Team/missless/internal/middleware"
"github.com/Two-Weeks-Team/missless/internal/store"
"github.com/Two-Weeks-Team/missless/internal/util"
)

Expand Down Expand Up @@ -51,11 +54,29 @@ func main() {
}
uploader := media.NewUploader(cfg.StorageBucket, storageClient)

// Firestore client (optional — nil if unavailable)
var firestoreClient *firestore.Client
if cfg.ProjectID != "" {
fc, err := firestore.NewClientWithDatabase(ctx, cfg.ProjectID, cfg.FirestoreDB)
if err != nil {
slog.Warn("firestore_client_init_failed", "error", err)
} else {
firestoreClient = fc
defer fc.Close()
}
}

// maxMemoriesPerPersona is the maximum number of memories stored per persona.
const maxMemoriesPerPersona = 100

sessionStore := store.NewFirestoreStore(cfg.ProjectID, firestoreClient)
memStore := memory.NewStore(maxMemoriesPerPersona, firestoreClient)

// HTTP mux
mux := http.NewServeMux()
sessions := auth.NewSessionStore()
handler.RegisterHealth(mux)
handler.RegisterWebSocket(mux, cfg, sessions)
handler.RegisterWebSocket(mux, cfg, sessions, sessionStore, memStore)
handler.RegisterOAuth(mux, cfg, sessions)
handler.RegisterUpload(mux, uploader)

Expand Down
32 changes: 28 additions & 4 deletions internal/handler/websocket.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,11 @@ import (
"github.com/Two-Weeks-Team/missless/internal/auth"
"github.com/Two-Weeks-Team/missless/internal/config"
"github.com/Two-Weeks-Team/missless/internal/live"
"github.com/Two-Weeks-Team/missless/internal/memory"
"github.com/Two-Weeks-Team/missless/internal/retry"
"github.com/Two-Weeks-Team/missless/internal/scene"
"github.com/Two-Weeks-Team/missless/internal/session"
"github.com/Two-Weeks-Team/missless/internal/store"
"github.com/gorilla/websocket"
"google.golang.org/genai"
)
Expand Down Expand Up @@ -58,13 +60,13 @@ func newUpgrader(cfg *config.Config) websocket.Upgrader {
}

// RegisterWebSocket registers the WebSocket endpoint for browser ↔ Go proxy.
func RegisterWebSocket(mux *http.ServeMux, cfg *config.Config, sessions *auth.SessionStore) {
func RegisterWebSocket(mux *http.ServeMux, cfg *config.Config, sessions *auth.SessionStore, sessionStore *store.FirestoreStore, memStore *memory.Store) {
mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
handleWebSocket(w, r, cfg, sessions)
handleWebSocket(w, r, cfg, sessions, sessionStore, memStore)
})
}

func handleWebSocket(w http.ResponseWriter, r *http.Request, cfg *config.Config, sessions *auth.SessionStore) {
func handleWebSocket(w http.ResponseWriter, r *http.Request, cfg *config.Config, sessions *auth.SessionStore, sessionStore *store.FirestoreStore, memStore *memory.Store) {
// Protection: Origin check (in upgrader) + rate limiter + connection limit.
// Session auth is NOT required here because the onboarding flow connects
// the WebSocket before the user completes OAuth login.
Expand All @@ -87,7 +89,12 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, cfg *config.Config,
activeWSConns.Add(1)
defer activeWSConns.Add(-1)

slog.Info("websocket_connected", "remote", r.RemoteAddr, "active_ws", activeWSConns.Load())
slog.Info("websocket_connected",
"remote", r.RemoteAddr,
"active_ws", activeWSConns.Load(),
"has_session_store", sessionStore != nil,
"has_memory_store", memStore != nil,
)

ctx, cancel := context.WithCancel(r.Context())
defer cancel()
Expand All @@ -114,6 +121,9 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, cfg *config.Config,
toolHandler := live.NewToolHandler()
toolHandler.SetGenerator(sceneGen)
toolHandler.SetGenaiClient(client)
if memStore != nil {
toolHandler.SetMemoryStore(memStore, "")

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Set persona ID when wiring memory store

recall_memory will never return results from the injected store because this initialization hard-codes an empty persona ID. In handleRecallMemory (internal/live/tools.go), the handler returns an empty memory list whenever pid == "", so any memories saved under real persona keys (for example from analysis) are unreachable at runtime even though memStore is present.

Useful? React with 👍 / 👎.

@gemini-code-assist gemini-code-assist Bot Feb 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The personaID passed to toolHandler.SetMemoryStore is an empty string. This will cause the recall_memory tool to be non-functional because the check for an empty personaID in handleRecallMemory will always pass, returning no results.

The personaID should be set on the toolHandler once it's known (likely after persona selection). This may require a mechanism for the session.Manager to update the toolHandler with the correct personaID when SetPersona is called.

}

// Connect to Live API with onboarding config and retry
liveConfig := mgr.BuildOnboardingConfig()
Expand Down Expand Up @@ -155,6 +165,20 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, cfg *config.Config,
cancel()
proxy.Close()

// Persist session state to Firestore on disconnect.
if sessionStore != nil {
saveCtx, saveCancel := context.WithTimeout(context.Background(), 3*time.Second)
defer saveCancel()
sd := &store.SessionData{
PersonaName: mgr.PersonaName(),
MatchedVoice: mgr.MatchedVoice(),
State: string(mgr.State()),
Comment on lines +173 to +175

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Populate session data from runtime state before persisting

handleWebSocket now writes PersonaName, MatchedVoice, and State from mgr on disconnect, but this handler never updates mgr after session.NewManager(...) (no SetPersona or state transitions are invoked in this code path), so persisted records will be saved as empty persona/voice and default onboarding state even after a real reunion flow. This makes the new Firestore persistence effectively incorrect for every session.

Useful? React with 👍 / 👎.

}
Comment on lines +172 to +176

@gemini-code-assist gemini-code-assist Bot Feb 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The session state being persisted to Firestore is incomplete. It's missing fields like MatchedVoice, ReunionCount, and LanguageCode which are part of the session state managed by session.Manager.

This could lead to inconsistent state if the session is restored later. Consider saving all relevant session data.

Note that languageCode is currently a private field in session.Manager and would need a getter to be accessed here.

		sd := &store.SessionData{
			PersonaName:  mgr.PersonaName(),
			State:        string(mgr.State()),
			MatchedVoice: mgr.MatchedVoice(),
			ReunionCount: mgr.ReunionCount(),
		}

if err := sessionStore.SaveSession(saveCtx, mgr.SessionID(), sd); err != nil {

@chatgpt-codex-connector chatgpt-codex-connector Bot Feb 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Use a stable session key when saving websocket state

The session is persisted under mgr.SessionID(), and in this handler that ID is initialized from r.RemoteAddr; source address/port values change across reconnects and can be reused, so saved state is not reliably tied to one logical user session and can be overwritten in multi-client or reconnect-heavy environments.

Useful? React with 👍 / 👎.

slog.Warn("session_persist_failed", "error", err)
}
}

shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer shutdownCancel()
mgr.Shutdown(shutdownCtx)
Expand Down
4 changes: 2 additions & 2 deletions internal/handler/websocket_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ func TestWebSocket_NoAuthRequired(t *testing.T) {

// Should NOT return 401 — proceeds to WebSocket upgrade (which fails in test
// since this isn't a real WS request, but the point is no auth rejection).
handleWebSocket(rec, req, cfg, sessions)
handleWebSocket(rec, req, cfg, sessions, nil, nil)

if rec.Code == http.StatusUnauthorized {
t.Fatal("WebSocket should not require session authentication")
Expand Down Expand Up @@ -158,7 +158,7 @@ func TestWebSocket_ConnectionLimit(t *testing.T) {
req := httptest.NewRequest("GET", "/ws", nil)
rec := httptest.NewRecorder()

handleWebSocket(rec, req, cfg, sessions)
handleWebSocket(rec, req, cfg, sessions, nil, nil)

if rec.Code != http.StatusServiceUnavailable {
t.Fatalf("expected 503 at connection limit, got %d", rec.Code)
Expand Down
Loading