Agent-System/internal/squad/workspace_test.go

package squad_test

import (
	"encoding/json"
	"os"
	"path/filepath"
	"testing"

	"chat-agent-system/internal/squad"
)

func TestWorkspaceEngine_PrepareWorkspace(t *testing.T) {
	tmpDir := t.TempDir()

	engine := squad.NewWorkspaceEngine(nil)

	req := squad.WorkspacePrepareRequest{
		WorkDir:      filepath.Join(tmpDir, "test_workspace_01"),
		SessionID:    squad.FlexSessionID("101"),
		ProjectID:    45,
		TaskKey:      "intake",
		StageTitle:   "用户初始需求收集表",
		ProjectName:  "内廷元无创神经刺激器",
		CustomerName: "漳州市圆山中医药产业研究院",
		DeviceType:   "二类医疗器械",
		IntendedUse:  "改善夜间打鼾与AHI",
		TargetMarket: "中国 (NMPA)",
		Backend:      "stub",
		Materials: []squad.MaterialItem{
			{
				Filename: "PIR.md",
				Content:  []byte("# PIR 需求表\n- 客户: 漳州市圆山中医药产业研究院\n- 重量: <=300g\n"),
			},
		},
		KnowledgeHits: []squad.KnowledgeHit{
			{
				DocCode: "GB 9706.1-2020",
				Title:   "医用电气设备安全",
				Content: "基本安全要求...",
				Source:  "知识库",
			},
		},
		TemplateContent: "# 用户初始需求收集表\n\n## 1. 客户与项目画像\n\n（四段式展开）\n",
		CustomAgents: []squad.AgentSpec{
			{
				Role:             "organizer",
				Name:             "材料整理专员",
				Model:            "claude-3-7-sonnet",
				Tools:            []string{"read", "write", "edit"},
				BashPermission:   "deny",
				Responsibilities: []string{"事实抽取", "缺口识别"},
				ForbiddenActions: []string{"严禁写价"},
			},
			{
				Role:             "writer",
				Name:             "文档撰写专员",
				Model:            "claude-3-7-sonnet",
				Tools:            []string{"read", "write", "edit"},
				BashPermission:   "deny",
				Responsibilities: []string{"撰写正文"},
				ForbiddenActions: []string{"严禁以省略写法跳过应写章节"},
			},
			{
				Role:             "leader",
				Name:             "商务对接负责人",
				Model:            "claude-3-7-sonnet",
				Tools:            []string{"read", "write", "edit"},
				BashPermission:   "deny",
				Responsibilities: []string{"二审与终审"},
				ForbiddenActions: []string{"严禁代做拍板"},
			},
		},
		WorkflowSteps: []squad.StepSpec{
			{
				ID:           "step_01_organize",
				Role:         "organizer",
				Title:        "抽取材料事实并回填字段",
				SubWorkspace: "tasks/01_organize",
				Mode:         "execute",
				Accept:       []string{"回填 fields.json", "事实与缺口写入 blackboard"},
			},
			{
				ID:           "step_02_draft",
				Role:         "writer",
				Title:        "撰写正文初稿",
				SubWorkspace: "tasks/02_draft",
				Mode:         "execute",
				Needs:        []string{"tasks/01_organize/output.md"},
				Accept:       []string{"按自描述模板四段式展开"},
			},
		},
	}

	ws, err := engine.PrepareWorkspace(req)
	if err != nil {
		t.Fatalf("PrepareWorkspace failed: %v", err)
	}

	if ws.WorkDir != req.WorkDir {
		t.Errorf("expected WorkDir %s, got %s", req.WorkDir, ws.WorkDir)
	}

	// 1. 校验 workspace.json
	wsFile := filepath.Join(ws.WorkDir, "workspace.json")
	if _, err := os.Stat(wsFile); err != nil {
		t.Fatalf("workspace.json does not exist: %v", err)
	}
	var wsData squad.WorkspaceMetadata
	wsBytes, _ := os.ReadFile(wsFile)
	if err := json.Unmarshal(wsBytes, &wsData); err != nil {
		t.Fatalf("unmarshal workspace.json failed: %v", err)
	}
	if wsData.TaskKey != "intake" {
		t.Errorf("expected TaskKey intake, got %s", wsData.TaskKey)
	}

	// 2. 校验 squad.json
	squadFile := filepath.Join(ws.WorkDir, "squad.json")
	if _, err := os.Stat(squadFile); err != nil {
		t.Fatalf("squad.json does not exist: %v", err)
	}
	var squadData squad.SquadConfig
	squadBytes, _ := os.ReadFile(squadFile)
	if err := json.Unmarshal(squadBytes, &squadData); err != nil {
		t.Fatalf("unmarshal squad.json failed: %v", err)
	}
	if len(squadData.Agents) != 3 {
		t.Errorf("expected 3 agents in squad.json, got %d", len(squadData.Agents))
	}

	// 3. 校验 workflow.json
	wfFile := filepath.Join(ws.WorkDir, "workflow.json")
	if _, err := os.Stat(wfFile); err != nil {
		t.Fatalf("workflow.json does not exist: %v", err)
	}
	var wfData squad.WorkflowPlan
	wfBytes, _ := os.ReadFile(wfFile)
	if err := json.Unmarshal(wfBytes, &wfData); err != nil {
		t.Fatalf("unmarshal workflow.json failed: %v", err)
	}
	if len(wfData.Steps) != 2 {
		t.Errorf("expected 2 steps in workflow.json, got %d", len(wfData.Steps))
	}

	// 4. 校验 context/ 目录与 context/index.json
	ctxIndexFile := filepath.Join(ws.WorkDir, "context", "index.json")
	if _, err := os.Stat(ctxIndexFile); err != nil {
		t.Fatalf("context/index.json does not exist: %v", err)
	}
	var ctxIndex squad.ContextIndex
	ctxBytes, _ := os.ReadFile(ctxIndexFile)
	if err := json.Unmarshal(ctxBytes, &ctxIndex); err != nil {
		t.Fatalf("unmarshal context/index.json failed: %v", err)
	}
	if len(ctxIndex.Items) == 0 {
		t.Errorf("expected non-empty context items in context/index.json")
	}

	// 5. 校验 blackboard.json
	bbFile := filepath.Join(ws.WorkDir, "blackboard.json")
	if _, err := os.Stat(bbFile); err != nil {
		t.Fatalf("blackboard.json does not exist: %v", err)
	}

	// 6. 校验子任务目录及 TASK.md / AGENTS.md
	task1Dir := filepath.Join(ws.WorkDir, "tasks", "01_organize")
	if _, err := os.Stat(task1Dir); err != nil {
		t.Fatalf("tasks/01_organize does not exist: %v", err)
	}
	if _, err := os.Stat(filepath.Join(task1Dir, "TASK.md")); err != nil {
		t.Fatalf("tasks/01_organize/TASK.md does not exist: %v", err)
	}
	if _, err := os.Stat(filepath.Join(task1Dir, "AGENTS.md")); err != nil {
		t.Fatalf("tasks/01_organize/AGENTS.md does not exist: %v", err)
	}
}