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

package runmgr

import (
	"context"
	"sync"
	"testing"

	agentllm "agent-native/runtime/llm"
	"agent-system-backend/internal/projectcontext"
	"agent-system-backend/pkg/extension"
	core "agent-system-core"
)

func TestProjectContextUsesPersistedRunAndRefreshesAuthorization(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunPending)
	allowed := true
	m.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return allowed, nil })
	called := 0
	m.SetProjectContextResolver(func(_ context.Context, req extension.ProjectContextRequest) (extension.ProjectContextBundle, error) {
		called++
		if req.Actor != string(run.Owner) || req.ProjectID != run.ProjectID || req.ConversationID != run.ConversationID || req.RunID != run.ID {
			t.Errorf("context scope did not come from persisted run: %+v", req)
		}
		return extension.ProjectContextBundle{Status: extension.ProjectContextReady, ProjectID: req.ProjectID}, nil
	})
	rc := &runCtx{mgr: m, run: run, ctx: context.Background()}
	if _, active, err := rc.ProjectContext(context.Background(), "work goal", 900); err != nil || !active || called != 1 {
		t.Fatalf("valid persisted scope rejected: active=%v called=%d err=%v", active, called, err)
	}
	allowed = false
	if _, _, err := rc.ProjectContext(context.Background(), "work goal", 900); err == nil || called != 1 {
		t.Fatal("revoked project still reached context provider")
	}
	allowed = true
	rc.run.ProjectID = "forged"
	if _, _, err := rc.ProjectContext(context.Background(), "work goal", 900); err == nil || called != 1 {
		t.Fatal("in-memory project override reached context provider")
	}
}

func TestProjectContextResumeReloadsCurrentVersionAndAuthorization(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunInterrupted)
	allowed, version := true, "v1"
	m.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return allowed, nil })
	m.SetProjectContextResolver(func(_ context.Context, req extension.ProjectContextRequest) (extension.ProjectContextBundle, error) {
		return extension.ProjectContextBundle{Status: extension.ProjectContextReady, ProjectID: req.ProjectID, Version: version}, nil
	})
	first := &runCtx{mgr: m, run: run, ctx: context.Background()}
	if b, _, err := first.ProjectContext(context.Background(), "work", 900); err != nil || b.Version != "v1" {
		t.Fatalf("initial context: %+v err=%v", b, err)
	}
	reloaded, err := st.GetRunByID(string(run.ID))
	if err != nil {
		t.Fatal(err)
	}
	resumed := &runCtx{mgr: m, run: reloaded, ctx: context.Background()}
	version = "v2"
	if b, _, err := resumed.ProjectContext(context.Background(), "work", 900); err != nil || b.Version != "v2" {
		t.Fatalf("resumed context reused stale version: %+v err=%v", b, err)
	}
	allowed = false
	if _, _, err := resumed.ProjectContext(context.Background(), "work", 900); err == nil {
		t.Fatal("resumed context ignored current revocation")
	}
}

func TestContextManifestPersistsByOwnerAndRun(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunPending)
	rc := &runCtx{mgr: m, run: run, ctx: context.Background()}
	manifest := extension.ProjectContextManifest{ProjectID: run.ProjectID, Version: "v3", Model: "model-1", InputSHA256: "abc", Included: []extension.ProjectContextManifestEntry{{SourceID: "s1", Revision: "r2", InputSHA256: "def"}}}
	if err := rc.SaveContextManifest(context.Background(), manifest); err != nil {
		t.Fatal(err)
	}
	got, err := st.ListContextManifests(string(run.Owner), string(run.ID))
	if err != nil || len(got) != 1 || got[0].Version != "v3" || got[0].Included[0].SourceID != "s1" {
		t.Fatalf("manifest missing from persisted run audit: got=%+v err=%v", got, err)
	}
	other, err := st.ListContextManifests("other", string(run.ID))
	if err != nil || len(other) != 0 {
		t.Fatalf("cross-owner manifest visible: got=%+v err=%v", other, err)
	}
}

func TestContextManifestRejectsDocumentDutyRevocationBeforeSend(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunPending)
	m.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	epoch := "duty-1"
	m.SetProjectContextResolver(func(_ context.Context, req extension.ProjectContextRequest) (extension.ProjectContextBundle, error) {
		return extension.ProjectContextBundle{Status: extension.ProjectContextReady, ProjectID: req.ProjectID, Version: "v1", AuthorizationEpoch: epoch,
			Items: []extension.ProjectContextItem{{SourceID: "restricted-doc", Text: "private text", Required: true}}}, nil
	})
	rc := &runCtx{mgr: m, run: run, ctx: context.Background()}
	bundle, active, err := rc.ProjectContext(context.Background(), "work", 900)
	if err != nil || !active {
		t.Fatalf("initial context: active=%v err=%v", active, err)
	}
	_, manifest, err := projectcontext.Assemble(bundle, 900)
	if err != nil {
		t.Fatal(err)
	}
	manifest.InputSHA256 = "request-hash"
	epoch = "duty-2" // project work remains allowed; document duty changed
	if err := rc.SaveContextManifest(context.Background(), manifest); err == nil {
		t.Fatal("stale document duty accepted at model send")
	}
	stored, err := st.ListContextManifests(string(run.Owner), string(run.ID))
	if err != nil || len(stored) != 0 {
		t.Fatalf("revoked context wrote a send receipt: %+v err=%v", stored, err)
	}
}

func TestProjectEvidenceOnlyReadsCurrentBundleReference(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunPending)
	m.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	listed := true
	m.SetProjectContextResolver(func(_ context.Context, req extension.ProjectContextRequest) (extension.ProjectContextBundle, error) {
		b := extension.ProjectContextBundle{Status: extension.ProjectContextReady, ProjectID: req.ProjectID}
		if listed {
			b.Evidence = []extension.ProjectEvidenceRef{{ID: "ref-1", SourceID: "doc-1", Revision: "r3", Hash: "source-hash"}}
		}
		return b, nil
	})
	read := 0
	m.SetProjectEvidenceReader(func(_ context.Context, req extension.ProjectEvidenceRequest) (extension.ProjectEvidence, error) {
		read++
		if req.Scope.RunID != run.ID || req.Reference.ID != "ref-1" {
			t.Errorf("untrusted evidence request: %+v", req)
		}
		return extension.ProjectEvidence{Reference: req.Reference, Text: "authorized detail"}, nil
	})
	rc := &runCtx{mgr: m, run: run, ctx: context.Background()}
	if _, err := rc.ProjectEvidence(context.Background(), "forged-ref"); err == nil || read != 0 {
		t.Fatal("arbitrary reference reached provider")
	}
	evidence, err := rc.ProjectEvidence(context.Background(), "ref-1")
	if err != nil || evidence.Text != "authorized detail" || read != 1 {
		t.Fatalf("current reference rejected: evidence=%+v read=%d err=%v", evidence, read, err)
	}
	if got := rc.ContextEvidenceInRequest([]agentllm.Message{{Role: agentllm.RoleTool, Content: "authorized detail"}}); len(got) != 1 || got[0].SourceID != "doc-1" {
		t.Fatalf("actual evidence read missing from request manifest: %+v", got)
	}
	if got := rc.ContextEvidenceInRequest([]agentllm.Message{{Role: agentllm.RoleTool, Content: "cleared old tool result"}}); len(got) != 0 {
		t.Fatalf("cleared evidence still counted as model input: %+v", got)
	}
	listed = false
	if _, err := rc.ProjectEvidence(context.Background(), "ref-1"); err == nil || read != 1 {
		t.Fatal("revoked reference reached provider")
	}
	if err := rc.SaveContextManifest(context.Background(), extension.ProjectContextManifest{ProjectID: run.ProjectID, InputSHA256: "new-request-hash"}); err == nil {
		t.Fatal("new model request accepted after an earlier evidence read was revoked")
	}
}

func TestConcurrentProjectEvidenceReadsKeepTrustedScope(t *testing.T) {
	st := openRecoveryStore(t)
	m := newRecoveryManager(t, st)
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunPending)
	m.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	ref := extension.ProjectEvidenceRef{ID: "r1", SourceID: "doc", Revision: "v1"}
	m.SetProjectContextResolver(func(_ context.Context, req extension.ProjectContextRequest) (extension.ProjectContextBundle, error) {
		return extension.ProjectContextBundle{Status: extension.ProjectContextReady, ProjectID: req.ProjectID, Evidence: []extension.ProjectEvidenceRef{ref}}, nil
	})
	m.SetProjectEvidenceReader(func(_ context.Context, req extension.ProjectEvidenceRequest) (extension.ProjectEvidence, error) {
		return extension.ProjectEvidence{Reference: req.Reference, Text: "body"}, nil
	})
	rc := &runCtx{mgr: m, run: run, ctx: context.Background()}
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if _, err := rc.ProjectEvidence(context.Background(), "r1"); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()
}