Agent-System/internal/taskrun/finalize_test.go

package taskrun_test

import (
	"context"
	"strings"
	"sync"
	"testing"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/taskrun"
)

func TestCommands_DuplicateFinalizeMemory(t *testing.T) {
	ctxMgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.TimeDecayConfig{})
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), "sess-dup", 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"task_status": "running"})
	if err != nil {
		t.Fatal(err)
	}
	cmds := &taskrun.Commands{Ctx: ctxMgr, Env: "test", StorageIdentity: "env:test"}
	rec := &taskrun.RunRecord{RunID: "run_dup_mem", SessionID: "sess-dup", TaskKey: "intake", AssistantMsgID: asstID}
	res := &bizdev.TaskRunResult{
		Success: true, TaskKey: "intake", ArtifactTitle: "用户初始需求收集表", MarkdownBody: "# body\n",
		DocRecord: &docgov.DocumentRecord{DocID: "doc_dup", ContentSHA256: strings.Repeat("a", 64), Version: 1},
	}
	st1, err := cmds.Finalize(context.Background(), rec, res)
	if err != nil {
		t.Fatal(err)
	}
	msg, err := ctxMgr.GetMessageByID(context.Background(), asstID)
	if err != nil {
		t.Fatal(err)
	}
	if msg.Content == "任务执行中…" {
		t.Fatal("placeholder not replaced")
	}
	st2, err := cmds.Finalize(context.Background(), rec, res)
	if err != nil {
		t.Fatal(err)
	}
	if st1 != "completed" || st2 != "completed" {
		t.Fatalf("status %s %s", st1, st2)
	}
	msg2, _ := ctxMgr.GetMessageByID(context.Background(), asstID)
	if msg2.ID != msg.ID || msg2.Content != msg.Content {
		t.Fatal("duplicate finalize mutated the assistant turn")
	}
}

func TestCommands_CancelWinsOverFinalizeMemory(t *testing.T) {
	ctxMgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.TimeDecayConfig{})
	asstID, _ := ctxMgr.RecordUserMessageWithID(context.Background(), "sess-race", 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", nil)
	cmds := &taskrun.Commands{Ctx: ctxMgr, Env: "test", StorageIdentity: "env:test"}
	rec := &taskrun.RunRecord{RunID: "run_race_mem", SessionID: "sess-race", AssistantMsgID: asstID}
	st, err := cmds.Cancel(context.Background(), rec)
	if err != nil {
		t.Fatal(err)
	}
	if st != taskrun.StatusCancelled {
		t.Fatalf("status=%s", st)
	}
	res := &bizdev.TaskRunResult{Success: true, MarkdownBody: "# later\n"}
	st2, err := cmds.Finalize(context.Background(), rec, res)
	if err != nil {
		t.Fatal(err)
	}
	if st2 != taskrun.StatusCancelled {
		t.Fatalf("finalize must adopt cancelled, got %s", st2)
	}
}

func TestCommands_ConcurrentMemorySingleTerminal(t *testing.T) {
	ctxMgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.TimeDecayConfig{})
	asstID, _ := ctxMgr.RecordUserMessageWithID(context.Background(), "sess-conc", 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", nil)
	cmds := &taskrun.Commands{Ctx: ctxMgr, Env: "test", StorageIdentity: "env:test"}
	rec := &taskrun.RunRecord{RunID: "run_conc_mem", SessionID: "sess-conc", AssistantMsgID: asstID}
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = cmds.Finalize(context.Background(), rec, &bizdev.TaskRunResult{Success: true, MarkdownBody: "# ok\n"})
	}()
	go func() {
		defer wg.Done()
		_, _ = cmds.Cancel(context.Background(), rec)
	}()
	wg.Wait()
	msg, err := ctxMgr.GetMessageByID(context.Background(), asstID)
	if err != nil {
		t.Fatal(err)
	}
	if msg.ID != asstID {
		t.Fatalf("assistant id changed %d", msg.ID)
	}
}

func TestRunManifest_CleanupGate(t *testing.T) {
	man := &taskrun.RunManifest{
		RunID: "run_clean", WorkspaceID: "ws", StorageIdentity: "env:dev",
		Members: []taskrun.ManifestMember{{LogicalID: "doc_a", Digest: "abc", Kind: taskrun.MemberKindDocument}},
	}
	if err := man.Freeze(); err != nil {
		t.Fatal(err)
	}
	if err := man.AssertCleanupAllowed(nil); err == nil {
		t.Fatal("cleanup without receipts must fail")
	}
	receipts := []taskrun.VerifiedReceipt{{
		LogicalID: "doc_a", Digest: "abc", StorageIdentity: "env:dev", Verified: true,
	}}
	if err := man.AssertCleanupAllowed(receipts); err != nil {
		t.Fatal(err)
	}
	bad := []taskrun.VerifiedReceipt{{
		LogicalID: "doc_a", Digest: "zzz", StorageIdentity: "env:dev", Verified: true,
	}}
	if err := man.AssertCleanupAllowed(bad); err == nil {
		t.Fatal("digest mismatch must fail")
	}
}

func TestCommands_RejectsSensitiveMarkdown(t *testing.T) {
	cmds := &taskrun.Commands{Env: "test", StorageIdentity: "env:test"}
	_, err := cmds.Finalize(context.Background(), &taskrun.RunRecord{RunID: "run_secret"}, &bizdev.TaskRunResult{
		Success: true, MarkdownBody: "api_key=super-secret",
	})
	if err == nil {
		t.Fatal("expected sensitive reject")
	}
}