diff --git a/server/fact_source_schema_test.go b/server/fact_source_schema_test.go new file mode 100644 index 0000000..b180ec9 --- /dev/null +++ b/server/fact_source_schema_test.go @@ -0,0 +1,108 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" +) + +// PG: on an existing database whose stored record_fact row predates the source +// field (and was edited by the operator), the startup seeders add the source +// enum without touching the operator's text, a second run changes nothing, and +// ToolResolve then serves the source property to the model. +func TestRecordFactSourceSchemaUpgrade(t *testing.T) { + skipWithoutDSN(t) + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { m.Close() }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + t.Cleanup(func() { agent.ToolResolve = nil }) + pg := m.pg + + old, err := pg.GetTool("record_fact") + if err != nil || old == nil { + t.Fatal("missing record_fact", err) + } + bindings, _ := json.Marshal(old.Agents) + t.Cleanup(func() { pg.UpdateTool(old.Key, old.Description, old.Schema, bindings, old.Enabled) }) + + // Simulate a pre-provenance row customised by the operator: no source + // property, an edited description and an edited parameter description. + var schema map[string]any + if err := json.Unmarshal(old.Schema, &schema); err != nil { + t.Fatal(err) + } + props := schema["properties"].(map[string]any) + delete(props, "source") + props["confidence"].(map[string]any)["description"] = "operator confidence text" + custom, _ := json.Marshal(schema) + if err := pg.UpdateTool("record_fact", "operator description", custom, bindings, true); err != nil { + t.Fatal(err) + } + if src := resolvedRecordFactProps(t)["source"]; src != nil { + t.Fatalf("precondition: stored schema should not expose source yet: %v", src) + } + + snapshot := func() (string, string, time.Time) { + var desc, sch string + var upd time.Time + if err := pg.QueryRow(`SELECT description, schema::text, updated_at FROM tools WHERE key='record_fact'`).Scan(&desc, &sch, &upd); err != nil { + t.Fatal(err) + } + return desc, sch, upd + } + s.seedOrchestrationTools() + desc1, sch1, upd1 := snapshot() + s.seedOrchestrationTools() + desc2, sch2, upd2 := snapshot() + if desc1 != desc2 || sch1 != sch2 || !upd1.Equal(upd2) { + t.Fatalf("second run changed the row: %q %s %v -> %q %s %v", desc1, sch1, upd1, desc2, sch2, upd2) + } + + tool, err := pg.GetTool("record_fact") + if err != nil || tool == nil { + t.Fatal(err) + } + var got struct { + Properties map[string]map[string]any `json:"properties"` + } + if err := json.Unmarshal(tool.Schema, &got); err != nil { + t.Fatal(err) + } + wantEnum := "[target_response tool_output third_party_intel operator_input inference]" + if src := got.Properties["source"]; src == nil || fmt.Sprint(src["enum"]) != wantEnum { + t.Fatalf("upgrade did not add the source enum: %s", tool.Schema) + } + if tool.Description != "operator description" || got.Properties["confidence"]["description"] != "operator confidence text" || len(got.Properties) != len(props)+1 { + t.Fatalf("operator configuration overwritten: %q %s", tool.Description, tool.Schema) + } + + resolved := resolvedRecordFactProps(t) + src, _ := resolved["source"].(map[string]any) + if src == nil || fmt.Sprint(src["enum"]) != wantEnum { + t.Fatalf("ToolResolve does not serve the source enum: %v", resolved) + } +} + +// resolvedRecordFactProps returns the properties of record_fact as ToolResolve +// serves it to the worker. +func resolvedRecordFactProps(t *testing.T) map[string]any { + t.Helper() + tools := agent.ToolResolve(context.Background(), "worker", agent.NewToolSet(nil, "").WorkerTools()) + for _, tl := range tools { + if tl.Name() == "record_fact" { + props, _ := tl.InputSchema()["properties"].(map[string]any) + return props + } + } + t.Fatal("worker lost record_fact") + return nil +} diff --git a/server/orchestration.go b/server/orchestration.go index d92ac00..2f184d9 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -571,11 +571,47 @@ func (s *Server) seedOrchestrationTools() { s.seedFindingWorkflowTools() s.seedFindingVerificationSchema() // S1: poe_level / evidence_quotes on update_finding_report, added once per DB s.seedFindingEvidenceSchema() // optional counterevidence / confidence_rationale on report_finding, added once per DB + s.seedFactSourceSchema() // source enum on record_fact, added once per DB // Note: pentest's default tool bindings need no migration — BuiltinToolSeeds already seeds // list_assets/insert_assets/report_finding/list_findings/list_companies together with // pentest at fresh initialization (no old DBs exist in this project, so no migration). } +// seedFactSourceSchema adds the source property (closed enum) to the record_fact +// schema of existing databases (SeedTool is first-insert only, and ToolResolve +// serves the stored schema, so without this the model would never see the enum). +// Idempotent: the property is added only when the stored schema lacks it, and user +// edits are preserved. +func (s *Server) seedFactSourceSchema() { + for _, seed := range agent.BuiltinToolSeeds() { + if seed.Key != "record_fact" { + continue + } + raw, err := json.Marshal(seed.Schema) + if err != nil { + log.Printf("[facts] record_fact schema: %v", err) + return + } + var obj map[string]json.RawMessage + if err := json.Unmarshal(raw, &obj); err != nil { + return + } + var properties map[string]json.RawMessage + if err := json.Unmarshal(obj["properties"], &properties); err != nil { + return + } + if len(properties["source"]) == 0 { + return + } + _, err = s.m.pg.Exec(`UPDATE tools SET schema=jsonb_set(schema,ARRAY['properties','source'],$1::jsonb,true),updated_at=now() +WHERE key='record_fact' AND system AND NOT(COALESCE(schema->'properties','{}'::jsonb) ? 'source')`, string(properties["source"])) + if err != nil { + log.Printf("[facts] upgrade record_fact.source: %v", err) + } + return + } +} + // refreshBuiltinToolSchemas propagates code schema/description changes on the // orchestration + platform tools into already-seeded rows ONCE per version flag — // SeedTool is first-insert-only, so a new param (e.g. spawn_task's llm_profile) never