feat(abac): enforce stage access on chat, search and SSE streams.

Tag messages with stage_key metadata, block sends for unauthorized roles, and filter history/search/case results by allowed stages.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
weichen
2026-09-13 02:50:23 +00:00
committed by weichen
co-authored by Cursor
parent 69a3be5463
commit d9f3d40423
13 changed files with 305 additions and 26 deletions
+13
View File
@@ -43,6 +43,7 @@ type Session struct {
ID string `json:"id"`
UserID int64 `json:"user_id"`
ProjectID int64 `json:"project_id"`
ActiveStageKey string `json:"active_stage_key,omitempty"`
UserCtx *coordinator.UserProjectContext `json:"user_ctx"`
Recognizer intent.Recognizer `json:"-"`
Coordinator *coordinator.CoordinatorEngine `json:"-"`
@@ -265,6 +266,12 @@ func (s *Session) sendMessageInternal(ctx context.Context, userInput string, str
}
if s.CtxMgr != nil {
userMeta := map[string]any{}
if s.ProjectID > 0 {
userMeta["project_id"] = s.ProjectID
}
if s.ActiveStageKey != "" {
userMeta["stage_key"] = s.ActiveStageKey
}
if clientTurnID != "" {
userMeta["client_turn_id"] = clientTurnID
userMeta["client_id"] = clientTurnID + ":user"
@@ -423,6 +430,12 @@ func (s *Session) sendMessageInternal(ctx context.Context, userInput string, str
asstMeta := map[string]any{
"reply_source": replySrc,
}
if s.ProjectID > 0 {
asstMeta["project_id"] = s.ProjectID
}
if s.ActiveStageKey != "" {
asstMeta["stage_key"] = s.ActiveStageKey
}
if recRes.TopIntent != nil {
asstMeta["turn_intent"] = string(recRes.TopIntent.Intent)
}
+24
View File
@@ -220,6 +220,30 @@ func (m *ContextManager) GetChatHistory(ctx context.Context, sessionID string, l
return res, nil
}
// GetSessionProjectID 返回会话关联项目 ID;无项目或未找到会话时返回 0。
func (m *ContextManager) GetSessionProjectID(ctx context.Context, sessionID string) (int64, error) {
if sessionID == "" {
return 0, errors.New("session_id required")
}
if m.pool != nil {
var projectID int64
err := m.pool.QueryRow(ctx, `
SELECT COALESCE(project_id, 0) FROM chat_sessions WHERE session_id = $1
`, sessionID).Scan(&projectID)
if err != nil {
return 0, fmt.Errorf("session not found: %w", err)
}
return projectID, nil
}
m.mu.RLock()
defer m.mu.RUnlock()
if s, ok := m.sessions[sessionID]; ok {
return s.ProjectID, nil
}
return 0, fmt.Errorf("session not found: %s", sessionID)
}
// GetSessionOwnerID 返回会话归属用户 ID;会话不存在时返回 error。
func (m *ContextManager) GetSessionOwnerID(ctx context.Context, sessionID string) (int64, error) {
if sessionID == "" {
+19 -9
View File
@@ -51,7 +51,7 @@ func (m *ContextManager) searchChatMessagesKeywordDB(ctx context.Context, q Chat
var builder strings.Builder
args := []any{q.SessionID, "%" + q.Query + "%"}
builder.WriteString(`
SELECT id, session_id, day_key::text, role, category, content
SELECT id, session_id, day_key::text, role, category, COALESCE(metadata->>'stage_key', ''), content
FROM chat_messages
WHERE session_id = $1 AND deleted_at IS NULL AND content ILIKE $2
`)
@@ -85,8 +85,8 @@ func (m *ContextManager) searchChatMessagesKeywordDB(ctx context.Context, q Chat
var hits []ChatSearchHit
for rows.Next() {
var id int64
var sid, dayKey, role, cat, content string
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &content); err != nil {
var sid, dayKey, role, cat, stageKey, content string
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content); err != nil {
return nil, err
}
hits = append(hits, ChatSearchHit{
@@ -95,6 +95,7 @@ func (m *ContextManager) searchChatMessagesKeywordDB(ctx context.Context, q Chat
DayKey: dayKey,
Role: role,
Category: cat,
StageKey: stageKey,
Excerpt: ExtractSearchExcerpt(content, q.Query, 40),
Score: 1.0,
MatchedBy: "keyword",
@@ -108,7 +109,7 @@ func (m *ContextManager) searchChatMessagesFuzzyDB(ctx context.Context, q ChatSe
var builder strings.Builder
args := []any{q.SessionID, q.Query}
builder.WriteString(`
SELECT id, session_id, day_key::text, role, category, content, similarity(content, $2) AS sim
SELECT id, session_id, day_key::text, role, category, COALESCE(metadata->>'stage_key', ''), content, similarity(content, $2) AS sim
FROM chat_messages
WHERE session_id = $1 AND deleted_at IS NULL AND similarity(content, $2) > 0.15
`)
@@ -142,9 +143,9 @@ func (m *ContextManager) searchChatMessagesFuzzyDB(ctx context.Context, q ChatSe
var hits []ChatSearchHit
for rows.Next() {
var id int64
var sid, dayKey, role, cat, content string
var sid, dayKey, role, cat, stageKey, content string
var sim float64
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &content, &sim); err != nil {
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content, &sim); err != nil {
return nil, err
}
hits = append(hits, ChatSearchHit{
@@ -153,6 +154,7 @@ func (m *ContextManager) searchChatMessagesFuzzyDB(ctx context.Context, q ChatSe
DayKey: dayKey,
Role: role,
Category: cat,
StageKey: stageKey,
Excerpt: ExtractSearchExcerpt(content, q.Query, 40),
Score: sim,
MatchedBy: "fuzzy",
@@ -176,7 +178,7 @@ func (m *ContextManager) searchChatMessagesSemanticDB(ctx context.Context, q Cha
var builder strings.Builder
args := []any{q.SessionID, vecStr}
builder.WriteString(`
SELECT m.id, m.session_id, m.day_key::text, m.role, m.category, m.content,
SELECT m.id, m.session_id, m.day_key::text, m.role, m.category, COALESCE(m.metadata->>'stage_key', ''), m.content,
(1 - (c.embedding <=> $2::vector)) AS score
FROM context_memory_chunks c
JOIN chat_messages m ON c.reference_id = 'msg:' || m.id
@@ -213,9 +215,9 @@ func (m *ContextManager) searchChatMessagesSemanticDB(ctx context.Context, q Cha
var hits []ChatSearchHit
for rows.Next() {
var id int64
var sid, dayKey, role, cat, content string
var sid, dayKey, role, cat, stageKey, content string
var score float64
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &content, &score); err != nil {
if err := rows.Scan(&id, &sid, &dayKey, &role, &cat, &stageKey, &content, &score); err != nil {
return nil, err
}
hits = append(hits, ChatSearchHit{
@@ -224,6 +226,7 @@ func (m *ContextManager) searchChatMessagesSemanticDB(ctx context.Context, q Cha
DayKey: dayKey,
Role: role,
Category: cat,
StageKey: stageKey,
Excerpt: ExtractSearchExcerpt(content, q.Query, 40),
Score: score,
MatchedBy: "semantic",
@@ -307,12 +310,19 @@ func (m *ContextManager) searchChatMessagesMemory(q ChatSearchQuery) ([]ChatSear
lowerContent := strings.ToLower(r.Content)
if strings.Contains(lowerContent, lowerQuery) {
stageKey := ""
if r.Metadata != nil {
if sk, ok := r.Metadata["stage_key"].(string); ok {
stageKey = sk
}
}
hits = append(hits, ChatSearchHit{
MessageID: r.ID,
SessionID: r.SessionID,
DayKey: r.DayKey,
Role: r.Role,
Category: r.Category,
StageKey: stageKey,
Excerpt: ExtractSearchExcerpt(r.Content, q.Query, 40),
Score: 1.0,
MatchedBy: "keyword",
+1
View File
@@ -288,6 +288,7 @@ type ChatSearchHit struct {
DayKey string `json:"day_key"`
Role string `json:"role"`
Category string `json:"category"`
StageKey string `json:"stage_key,omitempty"`
Excerpt string `json:"excerpt"`
Score float64 `json:"score"`
MatchedBy string `json:"matched_by"`
+19
View File
@@ -89,6 +89,25 @@ func (a *StageAccess) TaskKeyAllowed(taskKey string) bool {
return a.AllowedTaskKeys[taskKey]
}
// StageKeyAllowed 判断 stage_key 是否在允许范围内。
func (a *StageAccess) StageKeyAllowed(stageKey string) bool {
if a == nil || !a.Restricted || stageKey == "" {
return true
}
return a.AllowedStageKeys[stageKey]
}
// MessageStageKey 从消息 metadata 提取阶段键。
func MessageStageKey(metadata map[string]any) string {
if metadata == nil {
return ""
}
if sk, ok := metadata["stage_key"].(string); ok {
return sk
}
return ""
}
// TaskKeyToStage 从实例快照解析 task 所属阶段。
func TaskKeyToStage(plan WorkflowPlan, taskKey string) string {
for _, st := range plan.Stages {
@@ -0,0 +1,29 @@
package projectinstance
import "testing"
func TestStageKeyAllowed(t *testing.T) {
access := &StageAccess{
Restricted: true,
ActiveStageKey: "req_planning",
AllowedStageKeys: map[string]bool{"biz_intake": true},
}
if !access.StageKeyAllowed("biz_intake") {
t.Fatal("expected biz_intake allowed")
}
if access.StageKeyAllowed("req_planning") {
t.Fatal("expected req_planning denied for BD on completed intake only")
}
if !access.StageKeyAllowed("") {
t.Fatal("empty stage should pass when unrestricted check applies")
}
}
func TestMessageStageKey(t *testing.T) {
if got := MessageStageKey(map[string]any{"stage_key": "biz_intake"}); got != "biz_intake" {
t.Fatalf("expected biz_intake, got %q", got)
}
if got := MessageStageKey(nil); got != "" {
t.Fatalf("expected empty, got %q", got)
}
}
@@ -48,7 +48,12 @@ func TestResolveStageAccess_ActiveStageDutyOnly(t *testing.T) {
}
creatorAccess, _ := svc.ResolveStageAccess(ctx, 42, 1, true)
if !creatorAccess.TaskKeyAllowed("intake") {
if !creatorAccess.StageKeyAllowed("biz_intake") {
t.Fatal("creator should access active stage")
}
outsiderAccess, _ := svc.ResolveStageAccess(ctx, 42, 9, false)
if outsiderAccess.StageKeyAllowed("biz_intake") {
t.Fatal("user without duty should not access active stage chat")
}
}
+10 -8
View File
@@ -17,6 +17,11 @@ func (s *Server) handleChatSearch(w http.ResponseWriter, r *http.Request) {
return
}
user, ok := s.requireAuthenticatedUser(w, r)
if !ok {
return
}
if s.svc == nil || s.svc.ContextManager() == nil {
http.Error(w, "context manager not initialized", http.StatusServiceUnavailable)
return
@@ -46,16 +51,9 @@ func (s *Server) handleChatSearch(w http.ResponseWriter, r *http.Request) {
}
}
var userID int64
if uStr := r.URL.Query().Get("user_id"); uStr != "" {
if val, err := strconv.ParseInt(uStr, 10, 64); err == nil {
userID = val
}
}
searchQuery := ctxmgr.ChatSearchQuery{
SessionID: sessionID,
UserID: userID,
UserID: user.ID,
Query: query,
Mode: mode,
Category: category,
@@ -70,6 +68,10 @@ func (s *Server) handleChatSearch(w http.ResponseWriter, r *http.Request) {
return
}
if projectID, pidErr := s.svc.ContextManager().GetSessionProjectID(r.Context(), sessionID); pidErr == nil && projectID > 0 {
hits = s.filterChatSearchHitsByStageAccess(r.Context(), projectID, user.ID, hits)
}
if hits == nil {
hits = []ctxmgr.ChatSearchHit{}
}
+18
View File
@@ -17,6 +17,11 @@ func (s *Server) handleGetChatWindow(w http.ResponseWriter, r *http.Request) {
return
}
user, ok := s.requireAuthenticatedUser(w, r)
if !ok {
return
}
sessionID := r.URL.Query().Get("session_id")
if sessionID == "" {
http.Error(w, "missing session_id query param", http.StatusBadRequest)
@@ -40,6 +45,10 @@ func (s *Server) handleGetChatWindow(w http.ResponseWriter, r *http.Request) {
return
}
if projectID, pidErr := s.svc.ContextManager().GetSessionProjectID(r.Context(), sessionID); pidErr == nil && projectID > 0 {
s.applyChatWindowStageFilter(r.Context(), projectID, user.ID, res)
}
writeJSONResponse(w, http.StatusOK, res)
}
@@ -51,6 +60,11 @@ func (s *Server) handleGetChatDayMessages(w http.ResponseWriter, r *http.Request
return
}
user, ok := s.requireAuthenticatedUser(w, r)
if !ok {
return
}
sessionID := r.URL.Query().Get("session_id")
if sessionID == "" {
http.Error(w, "missing session_id query param", http.StatusBadRequest)
@@ -87,6 +101,10 @@ func (s *Server) handleGetChatDayMessages(w http.ResponseWriter, r *http.Request
return
}
if projectID, pidErr := s.svc.ContextManager().GetSessionProjectID(r.Context(), sessionID); pidErr == nil && projectID > 0 {
res.Messages = s.filterChatMessagesByStageAccess(r.Context(), projectID, user.ID, res.Messages)
}
writeJSONResponse(w, http.StatusOK, res)
}
+29
View File
@@ -142,6 +142,11 @@ func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
if req.ProjectID > 0 {
sess.ProjectID = req.ProjectID
}
s.enrichSessionWorkflowStage(r.Context(), sess, user.ID)
if sess.ProjectID > 0 && !s.allowActiveStageChatAccess(r.Context(), sess.ProjectID, user.ID) {
s.writeWorkflowAccessDenied(w)
return
}
resp, err := sess.SendMessage(r.Context(), req.Message)
if err != nil {
@@ -850,22 +855,36 @@ func (s *Server) handleGetChatMessages(w http.ResponseWriter, r *http.Request) {
offset = 0
}
user, ok := s.requireAuthenticatedUser(w, r)
if !ok {
return
}
msgs, err := s.svc.ContextManager().GetChatHistory(r.Context(), sessionID, limit, offset)
if err != nil {
http.Error(w, "get chat messages failed: "+err.Error(), http.StatusInternalServerError)
return
}
if projectID, pidErr := s.svc.ContextManager().GetSessionProjectID(r.Context(), sessionID); pidErr == nil && projectID > 0 {
msgs = s.filterChatMessagesByStageAccess(r.Context(), projectID, user.ID, msgs)
}
writeJSONResponse(w, http.StatusOK, msgs)
}
// handleRecommendCases 基于语义向量检索往期相似标杆案例
func (s *Server) handleRecommendCases(w http.ResponseWriter, r *http.Request) {
user, ok := s.requireAuthenticatedUser(w, r)
if !ok {
return
}
if s.svc == nil || s.svc.ContextManager() == nil {
http.Error(w, "context manager not initialized", http.StatusServiceUnavailable)
return
}
projectID, _ := strconv.ParseInt(r.URL.Query().Get("project_id"), 10, 64)
query := r.URL.Query().Get("query")
if query == "" {
http.Error(w, "missing query param", http.StatusBadRequest)
@@ -898,6 +917,16 @@ func (s *Server) handleRecommendCases(w http.ResponseWriter, r *http.Request) {
return
}
if projectID > 0 && len(cases) > 0 {
filtered := make([]ctxmgr.CaseRecommendationItem, 0, len(cases))
for _, c := range cases {
if s.allowWorkflowDocTaskAccess(r.Context(), projectID, user.ID, c.TaskKey) {
filtered = append(filtered, c)
}
}
cases = filtered
}
writeJSONResponse(w, http.StatusOK, cases)
}
+5
View File
@@ -179,6 +179,11 @@ func (s *Server) handleChatStream(w http.ResponseWriter, r *http.Request) {
if req.ProjectID > 0 {
sess.ProjectID = req.ProjectID
}
s.enrichSessionWorkflowStage(ctx, sess, req.UserID)
if sess.ProjectID > 0 && !s.allowActiveStageChatAccess(ctx, sess.ProjectID, req.UserID) {
writeWorkflowAccessDeniedSSE(w, flusher, req.SessionID)
return
}
if sess.UserCtx != nil {
if req.ProjectName != "" && (sess.UserCtx.Project.Name == "" || sess.UserCtx.Project.Name == "医疗器械研发项目" || sess.UserCtx.Project.Name == "当前项目") {
sess.UserCtx.Project.Name = req.ProjectName
+130 -7
View File
@@ -2,32 +2,63 @@ package server
import (
"context"
"encoding/json"
"net/http"
"chat-agent-system/internal/chat"
"chat-agent-system/internal/ctxmgr"
"chat-agent-system/internal/docgov"
"chat-agent-system/internal/projectinstance"
"chat-agent-system/internal/stream"
)
func (s *Server) allowWorkflowDocTaskAccess(ctx context.Context, projectID, userID int64, taskKey string) bool {
if taskKey == "" || s.instSvc == nil || !s.instSvc.Available() {
return true
func (s *Server) resolveStageAccess(ctx context.Context, projectID, userID int64) (*projectinstance.StageAccess, error) {
if s == nil || s.instSvc == nil || !s.instSvc.Available() || projectID <= 0 || userID <= 0 {
return &projectinstance.StageAccess{}, nil
}
isCreator, _ := s.pwSvc.IsProjectCreator(ctx, projectID, userID)
access, err := s.instSvc.ResolveStageAccess(ctx, projectID, userID, isCreator)
return s.instSvc.ResolveStageAccess(ctx, projectID, userID, isCreator)
}
func (s *Server) allowWorkflowDocTaskAccess(ctx context.Context, projectID, userID int64, taskKey string) bool {
if taskKey == "" {
return true
}
access, err := s.resolveStageAccess(ctx, projectID, userID)
if err != nil || access == nil || !access.Restricted {
return true
}
return access.TaskKeyAllowed(taskKey)
}
func (s *Server) allowActiveStageChatAccess(ctx context.Context, projectID, userID int64) bool {
access, err := s.resolveStageAccess(ctx, projectID, userID)
if err != nil || access == nil || !access.Restricted {
return true
}
return access.StageKeyAllowed(access.ActiveStageKey)
}
func (s *Server) enrichSessionWorkflowStage(ctx context.Context, sess *chat.Session, userID int64) {
if sess == nil || sess.ProjectID <= 0 || s.instSvc == nil || !s.instSvc.Available() {
return
}
settings, err := s.instSvc.GetSettings(ctx, sess.ProjectID, userID, false)
if err != nil || settings == nil {
return
}
sess.ActiveStageKey = settings.ActiveStageKey
}
func (s *Server) filterWorkflowDocsByStageAccess(
ctx context.Context,
projectID, userID int64,
docs []docgov.WorkflowDocumentSummary,
) []docgov.WorkflowDocumentSummary {
if len(docs) == 0 || s.instSvc == nil || !s.instSvc.Available() {
if len(docs) == 0 {
return docs
}
isCreator, _ := s.pwSvc.IsProjectCreator(ctx, projectID, userID)
access, err := s.instSvc.ResolveStageAccess(ctx, projectID, userID, isCreator)
access, err := s.resolveStageAccess(ctx, projectID, userID)
if err != nil || access == nil || !access.Restricted {
return docs
}
@@ -39,3 +70,95 @@ func (s *Server) filterWorkflowDocsByStageAccess(
}
return out
}
func (s *Server) filterChatMessagesByStageAccess(
ctx context.Context,
projectID, userID int64,
msgs []ctxmgr.ChatMessageRecord,
) []ctxmgr.ChatMessageRecord {
if len(msgs) == 0 || projectID <= 0 {
return msgs
}
access, err := s.resolveStageAccess(ctx, projectID, userID)
if err != nil || access == nil || !access.Restricted {
return msgs
}
out := make([]ctxmgr.ChatMessageRecord, 0, len(msgs))
for _, msg := range msgs {
stageKey := projectinstance.MessageStageKey(msg.Metadata)
if stageKey == "" {
continue
}
if access.StageKeyAllowed(stageKey) {
out = append(out, msg)
}
}
return out
}
func (s *Server) filterChatSearchHitsByStageAccess(
ctx context.Context,
projectID, userID int64,
hits []ctxmgr.ChatSearchHit,
) []ctxmgr.ChatSearchHit {
if len(hits) == 0 || projectID <= 0 {
return hits
}
access, err := s.resolveStageAccess(ctx, projectID, userID)
if err != nil || access == nil || !access.Restricted {
return hits
}
out := make([]ctxmgr.ChatSearchHit, 0, len(hits))
for _, hit := range hits {
if hit.StageKey == "" {
continue
}
if access.StageKeyAllowed(hit.StageKey) {
out = append(out, hit)
}
}
return out
}
func (s *Server) applyChatWindowStageFilter(ctx context.Context, projectID, userID int64, res *ctxmgr.ChatWindowResult) {
if res == nil || projectID <= 0 {
return
}
for i := range res.Days {
res.Days[i].Messages = s.filterChatMessagesByStageAccess(ctx, projectID, userID, res.Days[i].Messages)
res.Days[i].LoadedCount = len(res.Days[i].Messages)
res.Days[i].Count = len(res.Days[i].Messages)
if len(res.Days[i].Messages) > 0 {
res.Days[i].FirstMessageID = res.Days[i].Messages[0].ID
res.Days[i].LastMessageID = res.Days[i].Messages[len(res.Days[i].Messages)-1].ID
} else {
res.Days[i].FirstMessageID = 0
res.Days[i].LastMessageID = 0
}
}
}
func workflowAccessDeniedJSON() string {
return `{"error":"WORKFLOW_ACCESS_DENIED","message":"无权访问该阶段内容"}`
}
func (s *Server) writeWorkflowAccessDenied(w http.ResponseWriter) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusForbidden)
_, _ = w.Write([]byte(workflowAccessDeniedJSON()))
}
func writeWorkflowAccessDeniedSSE(w http.ResponseWriter, flusher http.Flusher, sessionID string) {
payload, _ := json.Marshal(stream.StreamEvent{
Type: stream.EventError,
ErrorInfo: &stream.ErrorInfo{
Code: "WORKFLOW_ACCESS_DENIED",
UserMessage: "无权访问该阶段对话内容",
SessionID: sessionID,
},
})
_, _ = w.Write([]byte("event: error\ndata: "))
_, _ = w.Write(payload)
_, _ = w.Write([]byte("\n\n"))
flusher.Flush()
}