fix: revert last change staging (#135)

Author: wjiayis

go.sum

@@ -55,8 +55,6 @@ github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
-github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE=
-github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4=
@@ -162,8 +160,6 @@ golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sU
 golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b h1:DXr+pvt3nC887026GRP39Ej11UATqWDmWuS99x26cD0=
 golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b/go.mod h1:4QTo5u+SEIbbKW1RacMZq1YEfOBqeXa19JeshGi+zc4=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
-golang.org/x/mod v0.27.0 h1:kb+q2PyFnEADO2IEF935ehFUXlWiNjJWtRNgBLSfbxQ=
-golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
@@ -192,8 +188,6 @@ golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
-golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg=
-golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
 gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=

internal/api/chat/create_conversation_message_stream_v2.go

@@ -281,7 +281,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 		APIKey: settings.OpenAIAPIKey,
 	}
 
-	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.UserID, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
+	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
 	if err != nil {
 		return s.sendStreamError(stream, err)
 	}
@@ -307,7 +307,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 			for i, bsonMsg := range conversation.InappChatHistory {
 				protoMessages[i] = mapper.BSONToChatMessageV2(bsonMsg)
 			}
-			title, err := s.aiClientV2.GetConversationTitleV2(ctx, conversation.UserID, protoMessages, llmProvider)
+			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider)
 			if err != nil {
 				s.logger.Error("Failed to get conversation title", "error", err, "conversationID", conversation.ID.Hex())
 				return

internal/api/grpc.go

@@ -15,7 +15,6 @@ import (
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
 	commentv1 "paperdebugger/pkg/gen/api/comment/v1"
 	projectv1 "paperdebugger/pkg/gen/api/project/v1"
-	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
 	userv1 "paperdebugger/pkg/gen/api/user/v1"
 
 	// "github.com/grpc-ecosystem/go-grpc-middleware"
@@ -107,7 +106,6 @@ func NewGrpcServer(
 	userServer userv1.UserServiceServer,
 	projectServer projectv1.ProjectServiceServer,
 	commentServer commentv1.CommentServiceServer,
-	usageServer usagev1.UsageServiceServer,
 ) *GrpcServer {
 	grpcServer := &GrpcServer{}
 	grpcServer.userService = userService
@@ -123,6 +121,5 @@ func NewGrpcServer(
 	userv1.RegisterUserServiceServer(grpcServer.Server, userServer)
 	projectv1.RegisterProjectServiceServer(grpcServer.Server, projectServer)
 	commentv1.RegisterCommentServiceServer(grpcServer.Server, commentServer)
-	usagev1.RegisterUsageServiceServer(grpcServer.Server, usageServer)
 	return grpcServer
 }

internal/api/server.go

@@ -6,23 +6,17 @@ import (
 	"fmt"
 	"net"
 	"net/http"
-	"os"
-	"os/signal"
 	"strings"
-	"syscall"
 
 	"paperdebugger/internal/libs/logger"
 	"paperdebugger/internal/libs/metadatautil"
 	"paperdebugger/internal/libs/shared"
-	"paperdebugger/internal/services"
-	aiclient "paperdebugger/internal/services/toolkit/client"
 	authv1 "paperdebugger/pkg/gen/api/auth/v1"
 	chatv1 "paperdebugger/pkg/gen/api/chat/v1"
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
 	commentv1 "paperdebugger/pkg/gen/api/comment/v1"
 	projectv1 "paperdebugger/pkg/gen/api/project/v1"
 	sharedv1 "paperdebugger/pkg/gen/api/shared/v1"
-	usagev1 "paperdebugger/pkg/gen/api/usage/v1"
 	userv1 "paperdebugger/pkg/gen/api/user/v1"
 
 	"github.com/gin-gonic/gin"
@@ -36,37 +30,25 @@ import (
 )
 
 type Server struct {
-	grpcServer     *GrpcServer
-	ginServer      *GinServer
-	pricingService *services.PricingService
-	aiClientV2     *aiclient.AIClientV2
+	grpcServer *GrpcServer
+	ginServer  *GinServer
 
 	logger *logger.Logger
 }
 
 func NewServer(
 	grpcServer *GrpcServer,
 	ginServer *GinServer,
-	pricingService *services.PricingService,
-	aiClientV2 *aiclient.AIClientV2,
 	logger *logger.Logger,
 ) *Server {
 	return &Server{
-		grpcServer:     grpcServer,
-		ginServer:      ginServer,
-		pricingService: pricingService,
-		aiClientV2:     aiClientV2,
-		logger:         logger,
+		grpcServer: grpcServer,
+		ginServer:  ginServer,
+		logger:     logger,
 	}
 }
 
 func (s *Server) Run(addr string) {
-	// Start the pricing updater in the background
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	s.pricingService.StartPriceUpdater(ctx)
-
 	listener, err := net.Listen("tcp", ":0")
 	if err != nil {
 		s.logger.Fatalf("failed to start grpc server listener: %v", err)
@@ -123,22 +105,6 @@ func (s *Server) Run(addr string) {
 		s.logger.Fatalf("failed to register comment service grpc gateway: %v", err)
 		return
 	}
-	err = usagev1.RegisterUsageServiceHandler(context.Background(), mux, client)
-	if err != nil {
-		s.logger.Fatalf("failed to register usage service grpc gateway: %v", err)
-		return
-	}
-
-	// Set up signal handling for graceful shutdown
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-
-	go func() {
-		<-sigChan
-		s.logger.Info("[PAPERDEBUGGER] received shutdown signal, shutting down gracefully...")
-		s.Shutdown()
-		os.Exit(0)
-	}()
 
 	s.logger.Infof("[PAPERDEBUGGER] http server listening on %s", addr)
 	s.ginServer.Any("/_pd/api/*path", func(c *gin.Context) { mux.ServeHTTP(c.Writer, c.Request) })
@@ -148,16 +114,6 @@ func (s *Server) Run(addr string) {
 	}
 }
 
-// Shutdown gracefully shuts down all server components.
-func (s *Server) Shutdown() {
-	s.logger.Info("[PAPERDEBUGGER] shutting down AI client (draining usage records)...")
-	s.aiClientV2.Shutdown()
-	s.logger.Info("[PAPERDEBUGGER] AI client shutdown complete")
-
-	s.grpcServer.GracefulStop()
-	s.logger.Info("[PAPERDEBUGGER] gRPC server shutdown complete")
-}
-
 func (s *Server) metadataAnnotator() func(ctx context.Context, req *http.Request) metadata.MD {
 	return func(ctx context.Context, req *http.Request) metadata.MD {
 		md := metadata.New(map[string]string{})

internal/api/usage/get_session_usage.go

@@ -6,7 +6,6 @@ import (
 
 	"paperdebugger/internal/libs/cfg"
 	"paperdebugger/internal/libs/logger"
-	"paperdebugger/internal/models"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
@@ -44,47 +43,5 @@ func NewDB(cfg *cfg.Cfg, logger *logger.Logger) (*DB, error) {
 	}
 
 	logger.Info("[MONGO] initialized")
-
-	db := &DB{Client: client, cfg: cfg, logger: logger}
-	db.ensureIndexes()
-	return db, nil
-}
-
-// ensureIndexes creates necessary indexes for the database collections.
-func (db *DB) ensureIndexes() {
-	sessions := db.Database("paperdebugger").Collection((models.LLMSession{}).CollectionName())
-
-	// TTL index: auto-delete sessions after 30 days past their expiry time
-	_, err := sessions.Indexes().CreateOne(context.Background(), mongo.IndexModel{
-		Keys:    bson.D{{Key: "session_expiry", Value: 1}},
-		Options: options.Index().SetExpireAfterSeconds(30 * 24 * 60 * 60),
-	})
-	if err != nil {
-		db.logger.Error("Failed to create TTL index on llm_sessions", "error", err)
-	}
-
-	// Compound index for efficient active session lookups
-	_, err = sessions.Indexes().CreateOne(context.Background(), mongo.IndexModel{
-		Keys: bson.D{
-			{Key: "user_id", Value: 1},
-			{Key: "session_expiry", Value: -1},
-		},
-	})
-	if err != nil {
-		db.logger.Error("Failed to create compound index on llm_sessions", "error", err)
-	}
-
-	// Unique compound index for session creation and queries.
-	// session_start is rounded to the second, so concurrent requests within the same
-	// second will conflict, triggering duplicate key handling in RecordUsage.
-	_, err = sessions.Indexes().CreateOne(context.Background(), mongo.IndexModel{
-		Keys: bson.D{
-			{Key: "user_id", Value: 1},
-			{Key: "session_start", Value: -1},
-		},
-		Options: options.Index().SetUnique(true),
-	})
-	if err != nil {
-		db.logger.Error("Failed to create session_start index on llm_sessions", "error", err)
-	}
+	return &DB{Client: client, cfg: cfg, logger: logger}, nil
 }