CyberStrikeAI/internal/handler/project_resume_test.go

package handler

import (
	"os"
	"path/filepath"
	"strings"
	"testing"

	"cyberstrike-ai/internal/config"
	"cyberstrike-ai/internal/database"
	"cyberstrike-ai/internal/project"
	"go.uber.org/zap"
)

func TestSessionContextReloadsProjectResultsEveryTurn(t *testing.T) {
	db, err := database.NewDB(filepath.Join(t.TempDir(), "context.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	p, err := db.CreateProject(&database.Project{Name: "shared"})
	if err != nil {
		t.Fatal(err)
	}
	cfg := &config.Config{Project: config.ProjectConfig{Enabled: true}}
	cfg.Agent.WorkspaceRootDir = t.TempDir()
	h := &AgentHandler{db: db, config: cfg, logger: zap.NewNop()}
	var ids []string
	for _, title := range []string{"yesterday", "today"} {
		c, err := db.CreateConversation(title, database.ConversationCreateMeta{ProjectID: p.ID})
		if err != nil {
			t.Fatal(err)
		}
		ids = append(ids, c.ID)
	}
	before := h.agentSessionContextBlock(ids[0])
	if strings.Contains(before, "parse_capture.py") {
		t.Fatal("unexpected old file")
	}
	root := project.WorkspaceRootDir(cfg.Agent.WorkspaceRootDir, p.ID, ids[0])
	if err := os.MkdirAll(filepath.Join(root, "artifacts"), 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(root, "artifacts", "parse_capture.py"), []byte("print('done')"), 0644); err != nil {
		t.Fatal(err)
	}
	if _, err := db.UpsertProjectFact(&database.ProjectFact{ProjectID: p.ID, FactKey: "artifact/parser", Category: "note", Summary: "verified parser", Body: "artifacts/parse_capture.py input capture.txt; output verified", Confidence: "confirmed"}); err != nil {
		t.Fatal(err)
	}
	v, err := db.CreateVulnerability(&database.Vulnerability{ProjectID: p.ID, Title: "tracked-finding", Severity: "low"})
	if err != nil {
		t.Fatal(err)
	}
	for _, id := range ids {
		block := h.agentSessionContextBlock(id)
		for _, want := range []string{"artifacts/parse_capture.py", "input capture.txt", v.ID, "tracked-finding", "先恢复成果"} {
			if !strings.Contains(block, want) {
				t.Errorf("conversation %s missing %s", id, want)
			}
		}
		if strings.Contains(block, "print('done')") {
			t.Fatal("script content injected without reading")
		}
	}
	// A separate project's workspace and findings must never be injected.
	other, err := db.CreateProject(&database.Project{Name: "other"})
	if err != nil {
		t.Fatal(err)
	}
	c, err := db.CreateConversation("isolated", database.ConversationCreateMeta{ProjectID: other.ID})
	if err != nil {
		t.Fatal(err)
	}
	isolated := h.agentSessionContextBlock(c.ID)
	if strings.Contains(isolated, "parse_capture.py") || strings.Contains(isolated, v.ID) {
		t.Fatal("project results leaked")
	}
}