First Commit

2026-10-09 08:38:16 +08:00
commit 0335d572de
756 changed files with 201663 additions and 0 deletions
@@ -0,0 +1,422 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/Autumn-27/artex/db"
+	actool "github.com/Autumn-27/norma/tool"
+)
+
+func callReadJSON(t *testing.T, tool actool.CoreTool, input string) any {
+	t.Helper()
+	result, err := tool.Call(context.Background(), json.RawMessage(input), nil)
+	if err != nil {
+		t.Fatalf("tool call: %v", err)
+	}
+	var out any
+	if err := json.Unmarshal([]byte(result.Flatten()), &out); err != nil {
+		t.Fatalf("decode tool result: %v; raw=%s", err, result.Flatten())
+	}
+	return out
+}
+
+func TestGraphOverviewExpandsAssociatedCompanyScope(t *testing.T) {
+	d := testDB(t)
+	defer d.Close()
+
+	companies := d.Companies()
+	companyID, _, err := companies.UpsertCompany(fmt.Sprintf("overview-scope-%d", time.Now().UnixNano()), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = companies.DeleteCompany(companyID) })
+
+	domain := fmt.Sprintf("overview-scope-%d.invalid", companyID)
+	ip := fmt.Sprintf("2001:db8:%x::42", companyID%0xffff)
+	cidr := fmt.Sprintf("2001:db8:%x:1::/64", companyID%0xffff)
+	icp := fmt.Sprintf("京 ICP 备 %d 号", companyID)
+	keyword := fmt.Sprintf("Scope Company %d", companyID)
+	inputs := []db.ScopeInput{
+		{Kind: "domain", Value: domain},
+		{Kind: "ip", Value: ip},
+		{Kind: "cidr", Value: cidr},
+		{Kind: "icp", Value: icp},
+		{Kind: "keyword", Value: keyword},
+	}
+	added, skipped, invalid, scopeErrors := companies.AddScopeInputs(companyID, inputs, "task context test")
+	if added != len(inputs) || skipped != 0 || invalid != 0 || len(scopeErrors) != 0 {
+		t.Fatalf("add company scope: added=%d skipped=%d invalid=%d errors=%v", added, skipped, invalid, scopeErrors)
+	}
+	assets := d.Assets()
+	assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: domain})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) })
+
+	task, err := d.CreateTaskWithOptions("company context", "read configured scope", db.TaskCreateOptions{
+		CompanyIDs: []int64{companyID},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
+
+	tools := NewToolSet(d.Exploration(task.ExplorationID), "planner")
+	tools.SetTaskID(task.ID)
+	tools.SetAssetStore(assets, companies)
+	linkedAssets, err := assets.QueryByTask(task.ID, "root_domain", 10, 0)
+	if err != nil || len(linkedAssets) != 1 || linkedAssets[0].ID != assetID {
+		t.Fatalf("company asset was not linked to task: assets=%+v err=%v", linkedAssets, err)
+	}
+	if linkedAssets[0].TaskSource != "company" || !strings.Contains(linkedAssets[0].TaskSourceSummary, companiesName(t, companies, companyID)) {
+		t.Fatalf("company asset provenance missing: %+v", linkedAssets[0])
+	}
+	overview := tools.graphOverviewData()
+	coverage, ok := overview["coverage"].(map[string]any)
+	if !ok {
+		t.Fatalf("coverage missing: %#v", overview["coverage"])
+	}
+	// graph_overview no longer flattens the task/company scope rows into the
+	// coverage block: upstream refactor 06a43f3 (slim down the situation-overview
+	// fields) dropped coverage["scope"], so company association now surfaces
+	// through host_count and the list_* asset tools below rather than a scope
+	// list. Assert the inherited company asset via host_count and the queries
+	// that follow.
+	if hc, _ := coverage["host_count"].(int); hc < 1 {
+		t.Fatalf("company asset host not counted in agent context: %#v", coverage["host_count"])
+	}
+	untested := callReadJSON(t, tools.listUntestedAssets(), `{"type":"root_domain","page":1,"page_size":10}`)
+	if !strings.Contains(fmt.Sprint(untested), domain) {
+		t.Fatalf("company asset missing from untested backlog: %#v", untested)
+	}
+	byTask := callReadJSON(t, tools.listAssets(), fmt.Sprintf(
+		`{"dsl":"task_id==%d","type":"root_domain","limit":10}`, task.ID,
+	))
+	if !strings.Contains(fmt.Sprint(byTask), domain) {
+		t.Fatalf("company asset missing from task-scoped agent query: %#v", byTask)
+	}
+	byCompany := callReadJSON(t, tools.listAssets(), fmt.Sprintf(
+		`{"dsl":"company_id==%d","type":"root_domain","limit":10}`, companyID,
+	))
+	if !strings.Contains(fmt.Sprint(byCompany), domain) {
+		t.Fatalf("company asset missing from company-scoped agent query: %#v", byCompany)
+	}
+	assetJSON, _ := json.Marshal(assetID)
+	intentID, err := tools.addOneIntent(intentItem{
+		Summary:  "test associated company asset",
+		AssetIDs: []json.RawMessage{assetJSON},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	workerAssets, err := assets.IntentAssets(task.ID)
+	if err != nil || len(workerAssets) != 1 || workerAssets[0].IntentID != intentID || workerAssets[0].AssetID != assetID {
+		t.Fatalf("worker target did not retain company asset: assets=%+v err=%v", workerAssets, err)
+	}
+
+	coverageDisabled := false
+	disabledTask, err := d.CreateTaskWithOptions("company context without coverage", "still expose company assets", db.TaskCreateOptions{
+		CompanyIDs: []int64{companyID}, CoverageEnabled: &coverageDisabled,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = d.DeleteTask(disabledTask.ID) })
+	disabledTools := NewToolSet(d.Exploration(disabledTask.ExplorationID), "planner")
+	disabledTools.SetTaskID(disabledTask.ID)
+	disabledTools.SetCoverageEnabled(false)
+	disabledTools.SetAssetStore(assets, companies)
+	disabledOverview := disabledTools.graphOverviewData()
+	disabledCoverage, ok := disabledOverview["coverage"].(map[string]any)
+	if !ok {
+		t.Fatalf("coverage-disabled task lost asset context: %#v", disabledOverview["coverage"])
+	}
+	if hc, _ := disabledCoverage["host_count"].(int); hc < 1 {
+		t.Fatalf("coverage-disabled task lost company asset host count: %#v", disabledCoverage["host_count"])
+	}
+	if _, exists := disabledCoverage["denominator"]; exists {
+		t.Fatalf("coverage-disabled task unexpectedly exposed metrics: %#v", disabledCoverage)
+	}
+	if linked, err := assets.QueryByTask(disabledTask.ID, "root_domain", 10, 0); err != nil || len(linked) != 1 || linked[0].ID != assetID {
+		t.Fatalf("coverage-disabled task asset link=%+v err=%v", linked, err)
+	}
+	disabledIntentID, err := disabledTools.addOneIntent(intentItem{
+		Summary:  "test associated company asset without coverage metrics",
+		AssetIDs: []json.RawMessage{assetJSON},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	disabledWorkerAssets, err := assets.IntentAssets(disabledTask.ID)
+	if err != nil || len(disabledWorkerAssets) != 1 || disabledWorkerAssets[0].IntentID != disabledIntentID || disabledWorkerAssets[0].AssetID != assetID {
+		t.Fatalf("coverage-disabled worker target=%+v err=%v", disabledWorkerAssets, err)
+	}
+}
+
+func companiesName(t *testing.T, companies *db.CompanyStore, companyID int64) string {
+	t.Helper()
+	company, err := companies.GetCompany(companyID)
+	if err != nil || company == nil {
+		t.Fatalf("company %d: company=%+v err=%v", companyID, company, err)
+	}
+	return company.Name
+}
+
+func TestBlackboardToolsReadDirectSources(t *testing.T) {
+	d := testDB(t)
+	defer d.Close()
+
+	grand, err := d.CreateTask("grand", "grand goal", nil, 0, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	source, err := d.CreateTaskWithOptions("source", "source goal", db.TaskCreateOptions{SourceTaskIDs: []int64{grand.ID}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	current, err := d.CreateTaskWithOptions("current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{source.ID}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		_ = d.DeleteTask(current.ID)
+		_ = d.DeleteTask(source.ID)
+		_ = d.DeleteTask(grand.ID)
+	})
+
+	grandStore := d.Exploration(grand.ExplorationID)
+	sourceStore := d.Exploration(source.ExplorationID)
+	currentStore := d.Exploration(current.ExplorationID)
+	grandFact, err := grandStore.AddNode(db.KindFact, map[string]any{"summary": "indirect-only"}, 0, "confirmed", "worker", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sourceFact, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "shared fact", "confidence": "verified"}, 0, "confirmed", "worker", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	sourceIntent, err := sourceStore.AddIntent(map[string]any{"summary": "shared work"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := sourceStore.SetIntentState(sourceIntent, "done"); err != nil {
+		t.Fatal(err)
+	}
+	sourceFinding, err := sourceStore.AddNode(db.KindFinding, map[string]any{
+		"summary": "shared finding", "vulnclass": "idor", "severity": "high",
+	}, 9, "confirmed", "worker", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := sourceStore.Link(sourceIntent, db.RelYields, sourceFinding); err != nil {
+		t.Fatal(err)
+	}
+	stepID, err := sourceStore.AppendActivity(db.Activity{
+		NodeID: &sourceIntent, Worker: "source-worker", Kind: "result", Tool: "HTTP",
+		Summary: "shared trace marker", Detail: "shared trace full detail",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for i := 0; i < 301; i++ {
+		if _, err := sourceStore.AddIntent(map[string]any{"summary": fmt.Sprintf("newer live source work %d", i)}, 1, nil, "planner"); err != nil {
+			t.Fatal(err)
+		}
+	}
+	assets := d.Assets()
+	host := fmt.Sprintf("overview-host-%d.invalid", source.ID)
+	assetID, err := assets.UpsertRootDomain(db.UpsertRootDomainReq{Domain: host, TaskID: source.ID})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _, _ = assets.DeleteByIDs([]int64{assetID}) })
+	if _, err := sourceStore.AddNode(db.KindFact, map[string]any{"summary": "source host anchor"}, 0, "confirmed", "worker", []int64{assetID}); err != nil {
+		t.Fatal(err)
+	}
+
+	tools := NewToolSet(currentStore, "worker")
+	tools.SetTaskID(current.ID)
+	tools.SetAssetStore(assets, assets.Companies())
+	parentID, _ := json.Marshal(sourceFact)
+	derivedID, err := tools.addOneIntent(intentItem{
+		Summary: "derive locally from shared fact", ParentIDs: []json.RawMessage{parentID},
+	})
+	if err != nil {
+		t.Fatalf("derive from inherited fact: %v", err)
+	}
+	if derived, err := currentStore.GetNode(derivedID); err != nil || derived == nil {
+		t.Fatalf("derived intent must be local: node=%+v err=%v", derived, err)
+	}
+	beforeFacts, _ := currentStore.ListByKind(db.KindFact, 100)
+	sourceIntentID, _ := json.Marshal(sourceIntent)
+	if _, err := tools.recordOneFact(factItem{Summary: "must not attach to inherited intent", IntentID: json.RawMessage(sourceIntentID)}, sourceIntent); err == nil {
+		t.Fatal("record_fact must reject inherited intent")
+	}
+	afterFacts, _ := currentStore.ListByKind(db.KindFact, 100)
+	if len(afterFacts) != len(beforeFacts) {
+		t.Fatalf("rejected inherited write persisted a fact: before=%d after=%d", len(beforeFacts), len(afterFacts))
+	}
+	overview := tools.graphOverviewData()
+	related, ok := overview["related_tasks"].([]map[string]any)
+	if !ok || len(related) != 1 {
+		t.Fatalf("related_tasks: %#v", overview["related_tasks"])
+	}
+	if related[0]["source_task_id"] != source.ID || related[0]["inherited"] != true {
+		t.Fatalf("related source provenance: %#v", related[0])
+	}
+	recentFacts, ok := related[0]["recent_facts"].([]map[string]any)
+	foundSourceFact := false
+	for _, fact := range recentFacts {
+		if fact["id"] == sourceFact {
+			foundSourceFact = true
+			break
+		}
+	}
+	if !ok || !foundSourceFact {
+		t.Fatalf("related fact summary: %#v", related[0]["recent_facts"])
+	}
+	if findings, ok := related[0]["recent_findings"].([]map[string]any); !ok || len(findings) == 0 || findings[0]["id"] != sourceFinding {
+		t.Fatalf("related finding summary: %#v", related[0]["recent_findings"])
+	}
+	results, ok := related[0]["recent_intent_results"].([]map[string]any)
+	if !ok || len(results) == 0 || results[0]["id"] != sourceIntent {
+		t.Fatalf("terminal source intent was starved by newer live work: %#v", related[0]["recent_intent_results"])
+	}
+	if results[0]["result_summary"] != "shared trace marker" {
+		t.Fatalf("terminal source intent result was not distilled: %#v", results[0])
+	}
+	coverage, ok := overview["coverage"].(map[string]any)
+	if !ok {
+		t.Fatalf("coverage missing from overview: %#v", overview["coverage"])
+	}
+	if coverage["host_count"] != 1 {
+		t.Fatalf("inherited host context missing (want host_count=1): %#v", coverage)
+	}
+
+	facts := callReadJSON(t, tools.listFacts(), `{}`).(map[string]any)["facts"].([]any)
+	seenSource, seenGrand := false, false
+	for _, raw := range facts {
+		item := raw.(map[string]any)
+		id := int64(item["id"].(float64))
+		if id == sourceFact {
+			seenSource = item["inherited"] == true && int64(item["source_task_id"].(float64)) == source.ID
+		}
+		seenGrand = seenGrand || id == grandFact
+	}
+	if !seenSource || seenGrand {
+		t.Fatalf("list_facts direct-only: source=%v grand=%v payload=%#v", seenSource, seenGrand, facts)
+	}
+
+	findings := callReadJSON(t, tools.listFindings(), `{}`).([]any)
+	if len(findings) != 1 {
+		t.Fatalf("list_findings: %#v", findings)
+	}
+	finding := findings[0].(map[string]any)
+	if finding["inherited"] != true || int64(finding["task_id"].(float64)) != source.ID || int64(finding["intent_id"].(float64)) != sourceIntent {
+		t.Fatalf("inherited finding provenance: %#v", finding)
+	}
+
+	detailInput, _ := json.Marshal(map[string]any{"id": sourceFact})
+	detail := callReadJSON(t, tools.nodeDetail(), string(detailInput)).(map[string]any)
+	if detail["inherited"] != true || int64(detail["source_task_id"].(float64)) != source.ID {
+		t.Fatalf("inherited node detail: %#v", detail)
+	}
+
+	traceInput, _ := json.Marshal(map[string]any{"intent_id": sourceIntent})
+	trace := callReadJSON(t, tools.getWorkerTrace(), string(traceInput)).(map[string]any)
+	if trace["inherited"] != true || int64(trace["source_task_id"].(float64)) != source.ID {
+		t.Fatalf("inherited trace: %#v", trace)
+	}
+	steps := trace["steps"].([]any)
+	if len(steps) != 1 || int64(steps[0].(map[string]any)["step_id"].(float64)) != stepID {
+		t.Fatalf("inherited trace steps: %#v", steps)
+	}
+	output := callReadJSON(t, tools.getWorkerOutput(), string(traceInput)).(map[string]any)
+	if output["inherited"] != true || output["final_text"] != "shared trace full detail" {
+		t.Fatalf("inherited worker output: %#v", output)
+	}
+
+	search := callReadJSON(t, tools.searchAllWorkerTraces(), `{"q":"shared trace marker"}`).(map[string]any)
+	hits := search["hits"].([]any)
+	if len(hits) != 1 || hits[0].(map[string]any)["inherited"] != true {
+		t.Fatalf("inherited trace search: %#v", hits)
+	}
+}
+
+// TestGetWorkerTraceStepIDsDegradeGracefully pins the over-cap behaviour: instead
+// of erroring, get_worker_trace returns the first 5 requested steps and tells the
+// model which ids it deferred, after de-duplicating and dropping invalid ids.
+func TestGetWorkerTraceStepIDsDegradeGracefully(t *testing.T) {
+	d := testDB(t)
+	defer d.Close()
+
+	task, err := d.CreateTask("trace-cap", "goal", nil, 0, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
+	store := d.Exploration(task.ExplorationID)
+	intent, err := store.AddIntent(map[string]any{"summary": "cap work"}, 1, nil, "planner")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var stepIDs []int64
+	for i := range 6 {
+		sid, err := store.AppendActivity(db.Activity{
+			NodeID: &intent, Worker: "w", Kind: "result", Tool: "HTTP",
+			Summary: fmt.Sprintf("step %d", i), Detail: fmt.Sprintf("detail %d", i),
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+		stepIDs = append(stepIDs, sid)
+	}
+
+	tools := NewToolSet(store, "worker")
+	tools.SetAssetStore(d.Assets(), d.Companies())
+
+	// Request 7 ids: a duplicate of the first, an invalid 0, then all 6 real ids.
+	// After dedup/cleanup that is 6 valid ids — one over the cap.
+	requested := []int64{stepIDs[0], stepIDs[0], 0, stepIDs[1], stepIDs[2], stepIDs[3], stepIDs[4], stepIDs[5]}
+	input, _ := json.Marshal(map[string]any{"intent_id": intent, "step_ids": requested})
+	res := callReadJSON(t, tools.getWorkerTrace(), string(input)).(map[string]any)
+
+	returned, _ := res["returned_step_ids"].([]any)
+	if len(returned) != 5 {
+		t.Fatalf("returned_step_ids=%v, want the first 5", res["returned_step_ids"])
+	}
+	// First 5 distinct valid ids, in request order.
+	wantReturned := []int64{stepIDs[0], stepIDs[1], stepIDs[2], stepIDs[3], stepIDs[4]}
+	for i, raw := range returned {
+		if int64(raw.(float64)) != wantReturned[i] {
+			t.Fatalf("returned[%d]=%v, want %d", i, raw, wantReturned[i])
+		}
+	}
+	omitted, _ := res["omitted_step_ids"].([]any)
+	if len(omitted) != 1 || int64(omitted[0].(float64)) != stepIDs[5] {
+		t.Fatalf("omitted_step_ids=%v, want [%d]", res["omitted_step_ids"], stepIDs[5])
+	}
+	if notice, _ := res["notice"].(string); notice == "" {
+		t.Fatalf("notice missing — model would not know a step was deferred")
+	}
+	if steps, _ := res["steps"].([]any); len(steps) != 5 {
+		t.Fatalf("steps=%d, want 5 detail rows", len(steps))
+	}
+
+	// At or under the cap: no notice, no omitted list.
+	okInput, _ := json.Marshal(map[string]any{"intent_id": intent, "step_ids": stepIDs[:3]})
+	okRes := callReadJSON(t, tools.getWorkerTrace(), string(okInput)).(map[string]any)
+	if _, hasNotice := okRes["notice"]; hasNotice {
+		t.Fatalf("notice present for an in-cap request: %#v", okRes["notice"])
+	}
+	if _, hasOmitted := okRes["omitted_step_ids"]; hasOmitted {
+		t.Fatalf("omitted_step_ids present for an in-cap request: %#v", okRes["omitted_step_ids"])
+	}
+}