Agent-System-Backend/internal/runmgr/knowledge_resource_test.go

package runmgr

import (
	"agent-system-backend/internal/rag"
	"agent-system-backend/internal/store"
	core "agent-system-core"
	"context"
	"encoding/json"
	"errors"
	"testing"
)

func TestRunKnowledgeResourceAuthorization(t *testing.T) {
	rc, st := newExecEnv(t)
	if err := st.AddKnowledge(core.KnowledgeDoc{ID: "allowed-slice", Owner: "bob", ProjectID: rc.run.ProjectID, Chunk: "resourcequery published"}, core.KnowledgeDoc{ID: "private-slice", Owner: "bob", ProjectID: rc.run.ProjectID, Chunk: "resourcequery private"}, core.KnowledgeDoc{ID: "restricted-self", Owner: "alice", ProjectID: rc.run.ProjectID, Chunk: "resourcequery self"}); err != nil {
		t.Fatal(err)
	}
	rc.mgr.SetRetriever(rag.NewPipeline(st, rc.mgr.settings, nil, nil))
	rc.mgr.SetSharedKnowledgeAllow(func(string, core.ID) bool { return true })
	rc.mgr.SetKnowledgeResourceAllow(func(actor string, d core.KnowledgeDoc) bool {
		return actor == "alice" && d.ID == "allowed-slice" && d.Owner == "bob"
	})
	hits, err := rc.Retrieve(core.RetrievalQuery{Text: "resourcequery"})
	if err != nil || len(hits) != 1 || hits[0].Doc.ID != "allowed-slice" {
		t.Fatalf("model-context resources: %+v %v", hits, err)
	}
}

func TestProjectWorkRevocationGuardsDirectExecution(t *testing.T) {
	rc, st := newExecEnv(t)
	rc.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return false, nil })
	conv, err := st.GetConversation("alice", rc.run.ConversationID)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := rc.mgr.StartConversationRun(conv, core.Message{}, nil, nil); !errors.Is(err, store.ErrNotFound) {
		t.Fatalf("direct run should reject: %v", err)
	}
	res := rc.ExecuteTool(core.ToolCallRequest{Name: "write_artifact", Arguments: json.RawMessage(`{"name":"forbidden.md","content":"x"}`)})
	if res.Status != "error" {
		t.Fatalf("revoked execution: %+v", res)
	}
	artifacts, err := st.ListArtifacts("alice", rc.run.ProjectID)
	if err != nil || len(artifacts) != 0 {
		t.Fatalf("revoked tool wrote: %+v %v", artifacts, err)
	}
}

func TestRunKnowledgeResolvesTrustedProjectScope(t *testing.T) {
	rc, st := newExecEnv(t)
	if err := st.AddKnowledge(core.KnowledgeDoc{ID: "slice-current", Owner: "bob", ProjectID: "domain-current", Chunk: "scopeneedle current"}, core.KnowledgeDoc{ID: "slice-other", Owner: "bob", ProjectID: "domain-other", Chunk: "scopeneedle other"}); err != nil {
		t.Fatal(err)
	}
	rc.mgr.SetRetriever(rag.NewPipeline(st, rc.mgr.settings, nil, nil))
	rc.mgr.SetKnowledgeResourceAllow(func(string, core.KnowledgeDoc) bool { return true })
	rc.mgr.SetKnowledgeProjectScope(func(_ context.Context, actor string, id core.ID) (core.ID, error) {
		if actor != "alice" || id != rc.run.ProjectID {
			t.Fatalf("untrusted scope actor=%s project=%s", actor, id)
		}
		return "domain-current", nil
	})
	hits, err := rc.Retrieve(core.RetrievalQuery{Text: "scopeneedle"})
	if err != nil || len(hits) != 1 || hits[0].Doc.ID != "slice-current" {
		t.Fatalf("model current scope: %+v %v", hits, err)
	}
}

func TestProjectWorkRevocationBlocksQueuedHITL(t *testing.T) {
	c, st := newExecEnv(t)
	c.run.ID = "revoked-run"
	c.run.Status = core.RunWaitingHITL
	if err := st.CreateRun(&c.run); err != nil {
		t.Fatal(err)
	}
	h := &runHandle{hitlCh: make(chan hitlAnswer, 1)}
	c.mgr.active = map[core.ID]*runHandle{c.run.ID: h}
	card := core.HITLCard{ID: "revoked-card", RunID: c.run.ID, ConversationID: c.run.ConversationID}
	if err := st.SaveHITLCard("alice", &card); err != nil {
		t.Fatal(err)
	}
	c.mgr.SetProjectWorkAuthorization(func(_ context.Context, owner string, id core.ID) (bool, error) {
		if owner != "alice" || id != c.run.ProjectID {
			t.Fatalf("untrusted resolution identity owner=%s project=%s", owner, id)
		}
		return false, nil
	})
	if _, _, err := c.mgr.ResolveHITL("alice", card.ID, core.HITLResponse{ChoiceID: "yes"}); !errors.Is(err, store.ErrNotFound) {
		t.Fatalf("revoked confirmation should reject: %v", err)
	}
	persisted, err := st.GetHITLCard("alice", card.ID)
	if err != nil || persisted.Resolved || len(h.hitlCh) != 0 {
		t.Fatalf("revoked confirmation changed gate: card=%+v channel=%d error=%v", persisted, len(h.hitlCh), err)
	}
	run, err := st.GetRunByID(c.run.ID)
	if err != nil || run.Status != core.RunWaitingHITL {
		t.Fatalf("revoked confirmation resumed run: %+v %v", run, err)
	}
	payload, _ := json.Marshal(HITLCommandPayload{Owner: "alice", CardID: card.ID, Kind: "yes"})
	c.mgr.applyHITLCommand(ownedCommand{CommandID: "queued-approval", Payload: string(payload)})
	if len(h.hitlCh) != 0 {
		t.Fatal("queued approval delivered a revoked answer")
	}
	run, err = st.GetRunByID(c.run.ID)
	if err != nil || run.Status != core.RunWaitingHITL {
		t.Fatalf("queued approval resumed run: %+v %v", run, err)
	}
}

func TestProjectWorkRevocationAfterToolConfirmation(t *testing.T) {
	for _, choice := range []string{"yes", "session"} {
		t.Run(choice, func(t *testing.T) {
			c, st := newExecEnv(t)
			c.h = &runHandle{hitlCh: make(chan hitlAnswer, 1)}
			c.h.hitlCh <- hitlAnswer{resp: core.HITLResponse{ChoiceID: choice}}
			c.prof.Tools.Confirm = []string{"write_artifact"}
			ready := true
			c.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { ok := ready; ready = false; return ok, nil })
			res := c.ExecuteTool(core.ToolCallRequest{CallID: "revoked-call", Name: "write_artifact", Arguments: json.RawMessage(`{"name":"revoked.txt","content":"x"}`)})
			artifacts, err := st.ListArtifacts("alice", c.run.ProjectID)
			if err != nil {
				t.Fatal(err)
			}
			if res.Status != "error" || len(artifacts) != 0 {
				t.Fatalf("tool executed after revoked confirmation: status=%s artifacts=%d", res.Status, len(artifacts))
			}
			if c.mgr.sessionAllows(c.run.ConversationID, "write_artifact") {
				t.Fatal("revoked confirmation granted session permission")
			}
		})
	}
}

func TestProjectWorkRevocationBlocksDirectArtifactWrite(t *testing.T) {
	c, st := newExecEnv(t)
	c.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return false, nil })
	if _, err := c.WriteArtifact("revoked.txt", "data", "text/plain", []byte("private")); !errors.Is(err, store.ErrNotFound) {
		t.Fatalf("direct artifact should reject: %v", err)
	}
	artifacts, err := st.ListArtifacts("alice", c.run.ProjectID)
	if err != nil || len(artifacts) != 0 {
		t.Fatalf("revoked artifact was written: %+v %v", artifacts, err)
	}
}