ARTEX-chixuweihu/server/credentials_test.go

package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"iter"
	"strings"
	"testing"
	"time"

	"github.com/Autumn-27/artex/db"
	"github.com/Autumn-27/artex/llmrec"
	"github.com/Autumn-27/norma/llm"
)

type credentialServerProvider struct {
	calls   int
	invalid bool
}

func (p *credentialServerProvider) Complete(_ context.Context, r llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
	p.calls++
	if p.invalid {
		return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("invalid fixture JSON")}}, "end_turn", llm.Usage{}, nil
	}
	return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock(`{"accounts":[{"username":"qa-user","password":"qa-pass-321"}]}`)}}, "end_turn", llm.Usage{InputTokens: 88, OutputTokens: 22}, nil
}
func (*credentialServerProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
	return func(func(llm.StreamEvent, error) bool) { panic("unexpected stream") }
}
func TestCredentialProcessingFlowBindingAndScopes(t *testing.T) {
	s, initial, request := trafficEvidenceServer(t)
	f, err := s.m.pg.GetFinding(initial.FindingID)
	if err != nil {
		t.Fatal(err)
	}
	taskID := *f.TaskID
	oldRecording := s.m.pg.GetBool(db.SettingTaskIntelligence, false)
	oldOn, oldBind := s.m.pg.GetBool(db.SettingCredentialAgent, false), s.m.pg.GetBool(db.SettingCredentialTraffic, false)
	defer func() {
		_ = s.m.pg.SetBool(db.SettingTaskIntelligence, oldRecording)
		_ = s.m.pg.SetCredentialAgentSettings(t.Context(), &oldOn, nil)
		if oldOn {
			_ = s.m.pg.SetCredentialAgentSettings(t.Context(), nil, &oldBind)
		}
	}()
	_ = s.m.pg.SetBool(db.SettingCredentialAgent, false)
	_ = s.m.pg.SetBool(db.SettingCredentialTraffic, false)
	if r := request("PUT", "/api/settings", `{"credential_traffic_binding":true}`); r.Code != 400 {
		t.Fatal("ungated binding", r.Code, r.Body)
	}
	if r := request("PUT", "/api/settings", `{"credential_agent_enabled":true}`); r.Code != 400 {
		t.Fatal("agent enabled without recording", r.Code)
	}
	if r := request("PUT", "/api/settings", `{"task_intelligence_enabled":true}`); r.Code != 200 {
		t.Fatal(r.Code, r.Body)
	}
	if r := request("PUT", "/api/settings", `{"credential_agent_enabled":true}`); r.Code != 200 {
		t.Fatal(r.Code, r.Body)
	}
	if s.settingsPayload()[db.SettingCredentialTraffic] != false {
		t.Fatal("binding default must be off")
	}
	assetID, err := s.m.Assets().UpsertHTTPService(db.UpsertHTTPServiceReq{URL: "http://evidence.local", TaskID: taskID})
	if err != nil {
		t.Fatal(err)
	}
	defer s.m.pg.Exec(`DELETE FROM assets WHERE id=$1`, assetID)
	in := db.IntelligenceInput{Kind: "credential", Title: "structured fixture", Content: `{"username":"qa-user","password":"qa-pass-321"}`, Evidence: "http://evidence.local local fixture", Confidence: "observed", AssetIDs: []int64{assetID}}
	item, err := s.m.pg.RecordIntelligence(t.Context(), taskID, 0, "fixture", in)
	if err != nil {
		t.Fatal(err)
	}
	mock := &credentialServerProvider{}
	s.llmOn = true
	s.llmProv = llmrec.Wrap(mock, s.m.pg, "fixture", "fixture", "", "", func() bool { return false })
	s.llmCfg.Model = "fixture"
	first, err := s.m.pg.ClaimCredentialJob(t.Context(), false, taskID)
	if err != nil || first == nil {
		t.Fatal(err)
	}
	s.processCredentialJob(t.Context(), first)
	got, err := s.m.pg.GetIntelligence(t.Context(), item.ID)
	if err != nil || got.Credentials.State != "ready" || len(got.Credentials.Accounts) != 1 || mock.calls != 0 || got.Credentials.TrafficState != "off" {
		t.Fatalf("rules: %+v %v calls=%d", got.Credentials, err, mock.calls)
	}
	seedServerEvidenceFlow(t, s, "credential-proof", []byte(`{"username":"qa-user","password":"qa-pass-321"}`))
	seedServerEvidenceFlow(t, s, "credential-wrong", []byte(`{"username":"other-user","password":"qa-pass-321"}`))
	seedServerEvidenceFlow(t, s, "credential-other-host", []byte(`{"username":"qa-user","password":"qa-pass-321"}`))
	// Direct fixture inserts bypass the capture writer's full-text indexing.
	for _, id := range []string{"credential-proof", "credential-wrong", "credential-other-host"} {
		if _, err = s.m.traffic.DB().Exec(`INSERT INTO ex_fts(rowid,content) SELECT rowid,? FROM exchanges WHERE id=?`, `{"username":"qa-user","password":"qa-pass-321"}`, id); err != nil {
			t.Fatal(err)
		}
	}
	_, err = s.m.traffic.DB().Exec(`UPDATE exchanges SET host='other.local',url='http://other.local/' WHERE id='credential-other-host'`)
	if err != nil {
		t.Fatal(err)
	}
	if r := request("PUT", "/api/settings", `{"credential_traffic_binding":true}`); r.Code != 200 {
		t.Fatal(r.Code, r.Body)
	}
	second, err := s.m.pg.ClaimCredentialJob(t.Context(), true, taskID)
	if err != nil || second == nil {
		t.Fatal("no late binding job", err)
	}
	s.processCredentialJob(t.Context(), second)
	got, _ = s.m.pg.GetIntelligence(t.Context(), item.ID)
	if got.Credentials.TrafficState != "ready" || len(got.Credentials.Traffic) != 1 || got.Credentials.Traffic[0].ID != "credential-proof" || mock.calls != 0 || got.Status != "unverified" {
		t.Fatalf("binding: %+v calls=%d", got.Credentials, mock.calls)
	}
	base := fmt.Sprintf("/api/tasks/%d/intelligence/%d/traffic/credential-proof", taskID, item.ID)
	r := request("GET", base, "")
	if r.Code != 200 || !strings.Contains(r.Body.String(), "qa-pass-321") {
		t.Fatal("flow detail", r.Code, r.Body)
	}
	other, err := s.m.pg.CreateTask("credential scope", "local", nil, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer s.m.pg.DeleteTask(other.ID)
	if r = request("GET", fmt.Sprintf("/api/tasks/%d/intelligence/%d/traffic/credential-proof", other.ID, item.ID), ""); r.Code != 404 {
		t.Fatal("unrelated task reads flow", r.Code)
	}
	if _, err = s.m.pg.Exec(`INSERT INTO task_relations(task_id,source_task_id) VALUES($1,$2)`, other.ID, taskID); err != nil {
		t.Fatal(err)
	}
	if r = request("GET", fmt.Sprintf("/api/tasks/%d/intelligence/%d/traffic/credential-proof", other.ID, item.ID), ""); r.Code != 200 {
		t.Fatal("inherited read failed", r.Code)
	}
	if r = request("POST", fmt.Sprintf("/api/tasks/%d/intelligence/%d/credentials/retry", other.ID, item.ID), ""); r.Code != 404 {
		t.Fatal("inherited write accepted", r.Code)
	}
	if r = request("GET", fmt.Sprintf("/api/assets/%d/intelligence/%d/traffic/credential-proof", assetID, item.ID), ""); r.Code != 200 {
		t.Fatal("asset read failed", r.Code, r.Body)
	}
	_, _ = s.m.traffic.DB().Exec(`DELETE FROM exchange_bodies WHERE id='credential-proof'`)
	_, _ = s.m.traffic.DB().Exec(`DELETE FROM exchanges WHERE id='credential-proof'`)
	if r = request("GET", base, ""); r.Code != 200 || !strings.Contains(r.Body.String(), `"unavailable":true`) || !strings.Contains(r.Body.String(), "qa-pass-321") {
		t.Fatal("saved excerpt lost", r.Code, r.Body)
	}
	// Reprocessing must preserve bindings and excerpts after live traffic deletion.
	if r = request("POST", fmt.Sprintf("/api/tasks/%d/intelligence/%d/credentials/retry", taskID, item.ID), ""); r.Code != 202 {
		t.Fatal(r.Code, r.Body)
	}
	repeat, err := s.m.pg.ClaimCredentialJob(t.Context(), true, taskID)
	if err != nil || repeat == nil {
		t.Fatal(err)
	}
	s.processCredentialJob(t.Context(), repeat)
	preserved, _ := s.m.pg.GetIntelligence(t.Context(), item.ID)
	if len(preserved.Credentials.Traffic) != 1 || preserved.Credentials.Traffic[0].ID != "credential-proof" {
		t.Fatal("retry lost saved evidence")
	}
	// A prose credential invokes the isolated model exactly once, with usage
	// attributed to the dedicated lane, even when body recording is disabled.
	in.Title = "prose fixture"
	in.Content = "qa-user 的密码为 qa-pass-321"
	prose, err := s.m.pg.RecordIntelligence(t.Context(), taskID, 0, "fixture", in)
	if err != nil {
		t.Fatal(err)
	}
	job, err := s.m.pg.ClaimCredentialJob(t.Context(), true, taskID)
	if err != nil || job == nil {
		t.Fatal(err)
	}
	s.processCredentialJob(t.Context(), job)
	got, _ = s.m.pg.GetIntelligence(t.Context(), prose.ID)
	if got.Credentials.Method != "agent" || mock.calls != 1 || got.Credentials.InputTokens != 88 {
		t.Fatalf("model: %+v calls=%d", got.Credentials, mock.calls)
	}
	var count int
	err = s.m.pg.QueryRow(`SELECT count(*) FROM llm_usage WHERE task_id=$1 AND worker=$2 AND input_tokens=88`, fmt.Sprint(taskID), db.CredentialAgentKey).Scan(&count)
	if err != nil || count != 1 {
		t.Fatal("usage not attributed", count, err)
	}
	if r = request("PUT", "/api/settings", `{"credential_agent_enabled":false}`); r.Code != 200 {
		t.Fatal(r.Code, r.Body)
	}
	if s.m.pg.GetBool(db.SettingCredentialTraffic, true) {
		t.Fatal("dependent switch not cleared")
	}
	if r = request("POST", fmt.Sprintf("/api/tasks/%d/intelligence/%d/credentials/retry", taskID, item.ID), ""); r.Code != 409 {
		t.Fatal("off retry accepted", r.Code)
	}
	list := request("GET", fmt.Sprintf("/api/tasks/%d/intelligence", taskID), "")
	var page db.IntelligencePage
	if err = json.Unmarshal(list.Body.Bytes(), &page); err != nil || len(page.Items) != 2 || page.CredentialAgentEnabled || page.Items[0].Credentials == nil {
		t.Fatal("off hid derived credentials", err)
	}
}

type blockingCredentialProvider struct{ entered chan struct{} }

func (p *blockingCredentialProvider) Complete(ctx context.Context, _ llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
	close(p.entered)
	<-ctx.Done()
	return llm.Message{}, "", llm.Usage{}, errors.New("cancelled fixture")
}
func (*blockingCredentialProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
	return func(func(llm.StreamEvent, error) bool) { panic("unexpected streaming") }
}
func TestCredentialDisableCancelsInFlightModel(t *testing.T) {
	s, f, request := trafficEvidenceServer(t)
	finding, err := s.m.pg.GetFinding(f.FindingID)
	if err != nil {
		t.Fatal(err)
	}
	taskID := *finding.TaskID
	oldRecording := s.m.pg.GetBool(db.SettingTaskIntelligence, false)
	oldOn, oldBinding := s.m.pg.GetBool(db.SettingCredentialAgent, false), s.m.pg.GetBool(db.SettingCredentialTraffic, false)
	defer func() {
		_ = s.m.pg.SetBool(db.SettingTaskIntelligence, oldRecording)
		_ = s.m.pg.SetCredentialAgentSettings(t.Context(), &oldOn, nil)
		if oldOn {
			_ = s.m.pg.SetCredentialAgentSettings(t.Context(), nil, &oldBinding)
		}
	}()
	on := true
	_ = s.m.pg.SetIntelligenceSettings(t.Context(), &on, &on, nil)
	item, err := s.m.pg.RecordIntelligence(t.Context(), taskID, 0, "fixture", db.IntelligenceInput{Kind: "credential", Title: "cancel fixture", Content: "qa-user 的密码为 qa-pass-321", Evidence: "local fixture", Confidence: "observed"})
	if err != nil {
		t.Fatal(err)
	}
	mock := &blockingCredentialProvider{entered: make(chan struct{})}
	s.llmOn = true
	s.llmProv = mock
	job, err := s.m.pg.ClaimCredentialJob(t.Context(), false, taskID)
	if err != nil || job == nil {
		t.Fatal(err)
	}
	done := make(chan struct{})
	go func() { s.processCredentialJob(t.Context(), job); close(done) }()
	select {
	case <-mock.entered:
	case <-time.After(5 * time.Second):
		t.Fatal("model never started")
	}
	if r := request("PUT", "/api/settings", `{"credential_agent_enabled":false}`); r.Code != 200 {
		t.Fatal(r.Code, r.Body)
	}
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("switch did not cancel model")
	}
	got, err := s.m.pg.GetIntelligence(t.Context(), item.ID)
	if err != nil || got.Credentials.State != "queued" || len(got.Credentials.Accounts) != 0 || got.Content != item.Content {
		t.Fatal("off published result or lost raw input", err)
	}
	if next, err := s.m.pg.ClaimCredentialJob(t.Context(), false, taskID); err != nil || next != nil {
		t.Fatal("off admitted a new run")
	}
}
func TestCredentialTokenMatchesWithoutUsername(t *testing.T) {
	a := db.CredentialAccount{Username: "qa-user", Password: "qa-pass-321", Token: "qa-token-123456"}
	if !credentialMatches(`{"token":"qa-token-123456"}`, a) {
		t.Fatal("token-only response not matched")
	}
	if !credentialMatches(`{"username":"qa-user","password":"qa-pass-321"}`, a) {
		t.Fatal("password request not matched when a token is also known")
	}
	if credentialMatches(`{"username":"other","password":"qa-pass-321"}`, a) {
		t.Fatal("other account incorrectly matched")
	}
}

func TestCredentialFailedModelCanRetry(t *testing.T) {
	s, f, request := trafficEvidenceServer(t)
	finding, err := s.m.pg.GetFinding(f.FindingID)
	if err != nil {
		t.Fatal(err)
	}
	taskID := *finding.TaskID
	oldRecord, oldAgent, oldBind := s.m.pg.GetBool(db.SettingTaskIntelligence, false), s.m.pg.GetBool(db.SettingCredentialAgent, false), s.m.pg.GetBool(db.SettingCredentialTraffic, false)
	defer s.m.pg.SetIntelligenceSettings(t.Context(), &oldRecord, &oldAgent, &oldBind)
	on, off := true, false
	_ = s.m.pg.SetIntelligenceSettings(t.Context(), &on, &on, &off)
	item, err := s.m.pg.RecordIntelligence(t.Context(), taskID, 0, "fixture", db.IntelligenceInput{Kind: "credential", Title: "retry fixture", Content: "qa-user 的密码为 qa-pass-321", Evidence: "local fixture", Confidence: "observed"})
	if err != nil {
		t.Fatal(err)
	}
	mock := &credentialServerProvider{invalid: true}
	s.llmOn = true
	s.llmProv = mock
	job, err := s.m.pg.ClaimCredentialJob(t.Context(), false, taskID)
	if err != nil || job == nil {
		t.Fatal(err)
	}
	s.processCredentialJob(t.Context(), job)
	got, _ := s.m.pg.GetIntelligence(t.Context(), item.ID)
	if got.Credentials.State != "error" || mock.calls != 1 {
		t.Fatal("failure not recorded")
	}
	if again, err := s.m.pg.ClaimCredentialJob(t.Context(), false, taskID); err != nil || again != nil {
		t.Fatal("failed model automatically retried")
	}
	mock.invalid = false
	if r := request("POST", fmt.Sprintf("/api/tasks/%d/intelligence/%d/credentials/retry", taskID, item.ID), ""); r.Code != 202 {
		t.Fatal(r.Code, r.Body)
	}
	job, err = s.m.pg.ClaimCredentialJob(t.Context(), false, taskID)
	if err != nil || job == nil {
		t.Fatal(err)
	}
	s.processCredentialJob(t.Context(), job)
	got, _ = s.m.pg.GetIntelligence(t.Context(), item.ID)
	if got.Credentials.State != "ready" || len(got.Credentials.Accounts) != 1 || mock.calls != 2 {
		t.Fatal("explicit retry did not call model again")
	}
}