cursor-byok/internal/backend/forwarder/codebase_index_handler.go

package forwarder

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"

	"connectrpc.com/connect"
	"github.com/google/uuid"

	"cursor/gen/aiserverv1"
)

func (service *Service) CreateExperimentalIndex(_ context.Context, req *connect.Request[aiserverv1.CreateExperimentalIndexRequest]) (*connect.Response[aiserverv1.CreateExperimentalIndexResponse], error) {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	record, err := store.EnsureIndexed(CodebaseIndexRecord{
		Repo:      strings.TrimSpace(req.Msg.GetRepo()),
		TargetDir: strings.TrimSpace(req.Msg.GetTargetDir()),
		Files:     append([]string(nil), req.Msg.GetFiles()...),
	})
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	return connect.NewResponse(&aiserverv1.CreateExperimentalIndexResponse{IndexId: record.IndexID}), nil
}

func (service *Service) ListExperimentalIndexFiles(_ context.Context, req *connect.Request[aiserverv1.ListExperimentalIndexFilesRequest]) (*connect.Response[aiserverv1.ListExperimentalIndexFilesResponse], error) {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	indexID := strings.TrimSpace(req.Msg.GetIndexId())
	files, err := store.ListFiles(indexID)
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	return connect.NewResponse(&aiserverv1.ListExperimentalIndexFilesResponse{
		IndexId: indexID,
		Files:   codebaseIndexFileData(indexID, files),
	}), nil
}

func (service *Service) ListenExperimentalIndex(ctx context.Context, req *connect.Request[aiserverv1.ListenExperimentalIndexRequest], stream *connect.ServerStream[aiserverv1.ListenExperimentalIndexResponse]) error {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return connect.NewError(connect.CodeInternal, err)
	}
	indexID := strings.TrimSpace(req.Msg.GetIndexId())
	if indexID == "" {
		return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("index_id is required"))
	}
	_, err = store.EnsureIndexed(CodebaseIndexRecord{IndexID: indexID})
	if err != nil {
		return connect.NewError(connect.CodeInternal, err)
	}
	subscription, err := store.Subscribe(indexID)
	if err != nil {
		return connect.NewError(connect.CodeInternal, err)
	}
	defer store.Unsubscribe(subscription)
	if err := stream.Send(&aiserverv1.ListenExperimentalIndexResponse{
		IndexId: indexID,
		Item: &aiserverv1.ListenExperimentalIndexResponse_Ready{
			Ready: &aiserverv1.ListenExperimentalIndexResponse_ReadyItem{
				IndexId: indexID,
				Request: &aiserverv1.ListenExperimentalIndexRequest{
					IndexId: indexID,
				},
			},
		},
	}); err != nil {
		return err
	}
	var lastEventID int64
	for {
		events, err := store.ListEvents(indexID, lastEventID)
		if err != nil {
			return connect.NewError(connect.CodeInternal, err)
		}
		for _, event := range events {
			if event.ID > lastEventID {
				lastEventID = event.ID
			}
			if event.Kind != codebaseIndexEventKindRegister {
				continue
			}
			if err := stream.Send(codebaseIndexRegisterEventResponse(indexID, event)); err != nil {
				return err
			}
		}
		select {
		case <-ctx.Done():
			return nil
		case <-subscription.Signal:
		}
	}
}

func (service *Service) RegisterFileToIndex(_ context.Context, req *connect.Request[aiserverv1.RegisterFileToIndexRequest]) (*connect.Response[aiserverv1.RequestReceivedResponse], error) {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	reqUUID := uuid.NewString()
	_, _, err = store.RegisterFileWithRequest(req.Msg.GetIndexId(), reqUUID, CodebaseIndexFileRecord{
		WorkspaceRelativePath: strings.TrimSpace(req.Msg.GetWorkspaceRelativePath()),
		ContentHash:           codebaseFileContentHash(strings.Join(req.Msg.GetContent(), "\n")),
		Stage:                 codebaseIndexFileStage,
	})
	if err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	return connect.NewResponse(&aiserverv1.RequestReceivedResponse{ReqUuid: reqUUID}), nil
}

func (service *Service) SetupIndexDependencies(_ context.Context, req *connect.Request[aiserverv1.SetupIndexDependenciesRequest]) (*connect.Response[aiserverv1.SetupIndexDependenciesResponse], error) {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	if err := store.MarkDependenciesReady(req.Msg.GetIndexId()); err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	return connect.NewResponse(&aiserverv1.SetupIndexDependenciesResponse{}), nil
}

func (service *Service) ComputeIndexTopoSort(_ context.Context, req *connect.Request[aiserverv1.ComputeIndexTopoSortRequest]) (*connect.Response[aiserverv1.ComputeIndexTopoSortResponse], error) {
	store, err := service.requireCodebaseIndexStore()
	if err != nil {
		return nil, connect.NewError(connect.CodeInternal, err)
	}
	if err := store.MarkTopoSortReady(req.Msg.GetIndexId()); err != nil {
		return nil, connect.NewError(connect.CodeInvalidArgument, err)
	}
	return connect.NewResponse(&aiserverv1.ComputeIndexTopoSortResponse{}), nil
}

func (service *Service) requireCodebaseIndexStore() (*CodebaseIndexStore, error) {
	if service == nil || service.codebaseIndexStore == nil {
		return nil, fmt.Errorf("codebase index store is not initialized")
	}
	return service.codebaseIndexStore, nil
}

func codebaseIndexFileData(indexID string, files []CodebaseIndexFileRecord) []*aiserverv1.IndexFileData {
	result := make([]*aiserverv1.IndexFileData, 0, len(files))
	for index, file := range files {
		result = append(result, codebaseIndexFileDatum(indexID, file, int32(index+1)))
	}
	return result
}

func codebaseIndexFileDatum(indexID string, file CodebaseIndexFileRecord, fallbackOrder int32) *aiserverv1.IndexFileData {
	order := file.Order
	if order == 0 {
		order = fallbackOrder
	}
	return &aiserverv1.IndexFileData{
		IndexId:               indexID,
		WorkspaceRelativePath: file.WorkspaceRelativePath,
		Stage:                 file.Stage,
		Order:                 order,
	}
}

func codebaseIndexRegisterEventResponse(indexID string, event CodebaseIndexEventRecord) *aiserverv1.ListenExperimentalIndexResponse {
	fileData := codebaseIndexFileDatum(indexID, event.File, 1)
	return &aiserverv1.ListenExperimentalIndexResponse{
		IndexId: indexID,
		Item: &aiserverv1.ListenExperimentalIndexResponse_Register{
			Register: &aiserverv1.ListenExperimentalIndexResponse_RegisterItem{
				ReqUuid: strings.TrimSpace(event.ReqUUID),
				Request: &aiserverv1.RegisterFileToIndexRequest{
					IndexId:               indexID,
					WorkspaceRelativePath: event.File.WorkspaceRelativePath,
				},
				Response: &aiserverv1.RegisterFileToIndexResponse{
					FileId:            codebaseIndexFileID(indexID, event.File.WorkspaceRelativePath),
					RootContextNodeId: codebaseIndexFileID(indexID, event.File.WorkspaceRelativePath),
					FileData:          fileData,
				},
			},
		},
	}
}

func codebaseIndexFileID(indexID string, workspaceRelativePath string) string {
	trimmedIndexID := strings.TrimSpace(indexID)
	trimmedPath := strings.TrimSpace(workspaceRelativePath)
	if trimmedIndexID == "" || trimmedPath == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(trimmedIndexID + "\x00" + trimmedPath))
	return "file_" + hex.EncodeToString(sum[:])[:24]
}