prompt_worker_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595
  1. package securityaudit
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "reflect"
  7. "strings"
  8. "sync"
  9. "testing"
  10. "time"
  11. "github.com/stretchr/testify/require"
  12. )
  13. type fixedClock struct{ now time.Time }
  14. func (c fixedClock) Now() time.Time { return c.now }
  15. type advancingClock struct {
  16. mu sync.Mutex
  17. now time.Time
  18. step time.Duration
  19. }
  20. func (c *advancingClock) Now() time.Time {
  21. c.mu.Lock()
  22. defer c.mu.Unlock()
  23. c.now = c.now.Add(c.step)
  24. return c.now
  25. }
  26. type fakeConfigStore struct {
  27. cfg ActiveConfig
  28. active bool
  29. }
  30. func (s *fakeConfigStore) Start(context.Context) error { return nil }
  31. func (s *fakeConfigStore) Shutdown(context.Context) error { return nil }
  32. func (s *fakeConfigStore) Active() (ActiveConfig, bool) { return cloneActiveConfig(s.cfg), s.active }
  33. func (s *fakeConfigStore) EffectiveMode() Mode {
  34. if s.BlockingActivationDegraded() {
  35. return ModeBlocking
  36. }
  37. if !s.active {
  38. return ModeOff
  39. }
  40. return s.cfg.EffectiveMode()
  41. }
  42. func (s *fakeConfigStore) BlockingActivationDegraded() bool { return false }
  43. func (s *fakeConfigStore) Public() (PublicConfig, error) { return PublicConfig{}, nil }
  44. func (s *fakeConfigStore) Save(context.Context, UpdateConfigRequest, int64) (PublicConfig, error) {
  45. return PublicConfig{}, nil
  46. }
  47. func (s *fakeConfigStore) RuntimeState() (int64, int64, *time.Time, string) {
  48. return s.cfg.ConfigVersion, s.cfg.ConfigVersion, nil, ""
  49. }
  50. func (s *fakeConfigStore) Encrypt(value string) (string, error) { return value, nil }
  51. func (s *fakeConfigStore) Decrypt(value string) (string, error) { return value, nil }
  52. type fakeJobRepository struct {
  53. mu sync.Mutex
  54. trace *[]string
  55. createJob *Job
  56. createErr error
  57. publishErr error
  58. refreshErr error
  59. completeErr error
  60. retryErr error
  61. failErr error
  62. createdSnapshot PromptSnapshot
  63. markedCode string
  64. completedResult *NormalizedResult
  65. completedStore bool
  66. completeCount int
  67. eventCount int
  68. retryAt time.Time
  69. retryCode string
  70. retried int
  71. failedCode string
  72. failed int
  73. refreshes int
  74. claimQueue []*Job
  75. recordBlockingCalls int
  76. recordBlockingSnapshot PromptSnapshot
  77. recordBlockingResult *NormalizedResult
  78. recordBlockingErr error
  79. }
  80. func (r *fakeJobRepository) record(value string) {
  81. if r.trace != nil {
  82. *r.trace = append(*r.trace, value)
  83. }
  84. }
  85. func (r *fakeJobRepository) CreateStagingWithCapacity(_ context.Context, snapshot PromptSnapshot, _ int64, _, _ int) (*Job, error) {
  86. r.mu.Lock()
  87. defer r.mu.Unlock()
  88. r.record("create_staging")
  89. r.createdSnapshot = snapshot
  90. if r.createErr != nil {
  91. return nil, r.createErr
  92. }
  93. if r.createJob == nil {
  94. r.createJob = &Job{ID: 1, Snapshot: snapshot}
  95. }
  96. return r.createJob, nil
  97. }
  98. func (r *fakeJobRepository) PublishQueued(context.Context, int64) error {
  99. r.mu.Lock()
  100. defer r.mu.Unlock()
  101. r.record("publish_queued")
  102. return r.publishErr
  103. }
  104. func (r *fakeJobRepository) MarkStagingFailed(_ context.Context, _ int64, code, _ string) error {
  105. r.mu.Lock()
  106. defer r.mu.Unlock()
  107. r.record("mark_staging_failed")
  108. r.markedCode = code
  109. return nil
  110. }
  111. func (r *fakeJobRepository) ClaimNextJob(context.Context, time.Time) (*Job, bool, error) {
  112. r.mu.Lock()
  113. defer r.mu.Unlock()
  114. if len(r.claimQueue) == 0 {
  115. return nil, false, nil
  116. }
  117. job := r.claimQueue[0]
  118. r.claimQueue = r.claimQueue[1:]
  119. return job, true, nil
  120. }
  121. func (r *fakeJobRepository) RefreshLease(context.Context, int64, int64, time.Time) error {
  122. r.mu.Lock()
  123. defer r.mu.Unlock()
  124. r.refreshes++
  125. return r.refreshErr
  126. }
  127. func (r *fakeJobRepository) Complete(_ context.Context, _ *Job, result *NormalizedResult, storePass bool) (*Event, error) {
  128. r.mu.Lock()
  129. defer r.mu.Unlock()
  130. r.completeCount++
  131. r.completedResult, r.completedStore = result, storePass
  132. if r.completeErr != nil {
  133. return nil, r.completeErr
  134. }
  135. if result.Decision == EventPass && !storePass {
  136. return nil, nil
  137. }
  138. r.eventCount++
  139. return &Event{ID: 99, Decision: result.Decision}, nil
  140. }
  141. func (r *fakeJobRepository) Retry(_ context.Context, _, _ int64, next time.Time, code, _ string) error {
  142. r.mu.Lock()
  143. defer r.mu.Unlock()
  144. r.retried++
  145. r.retryAt, r.retryCode = next, code
  146. return r.retryErr
  147. }
  148. func (r *fakeJobRepository) Fail(_ context.Context, _, _ int64, code, _ string) error {
  149. r.mu.Lock()
  150. defer r.mu.Unlock()
  151. r.failed++
  152. r.failedCode = code
  153. return r.failErr
  154. }
  155. func (r *fakeJobRepository) ReclaimStale(context.Context, time.Time, time.Time, int) (int64, error) {
  156. return 0, nil
  157. }
  158. func (r *fakeJobRepository) QueueStats(context.Context) (QueueStats, error) { return QueueStats{}, nil }
  159. func (r *fakeJobRepository) RecordBlocking(_ context.Context, snapshot PromptSnapshot, _ int64, result *NormalizedResult, _ bool) (*Event, error) {
  160. r.mu.Lock()
  161. defer r.mu.Unlock()
  162. r.recordBlockingCalls++
  163. r.recordBlockingSnapshot, r.recordBlockingResult = snapshot, result
  164. return nil, r.recordBlockingErr
  165. }
  166. type fakePayloadStore struct {
  167. mu sync.Mutex
  168. trace *[]string
  169. values map[int64]string
  170. setErr error
  171. getErr error
  172. deleteErr error
  173. pingErr error
  174. setTTL time.Duration
  175. deleted []int64
  176. }
  177. func (s *fakePayloadStore) Set(_ context.Context, jobID int64, value string, ttl time.Duration) error {
  178. s.mu.Lock()
  179. defer s.mu.Unlock()
  180. if s.trace != nil {
  181. *s.trace = append(*s.trace, "payload_set")
  182. }
  183. if s.setErr != nil {
  184. return s.setErr
  185. }
  186. if s.values == nil {
  187. s.values = map[int64]string{}
  188. }
  189. s.values[jobID], s.setTTL = value, ttl
  190. return nil
  191. }
  192. func (s *fakePayloadStore) Get(_ context.Context, jobID int64) (string, error) {
  193. s.mu.Lock()
  194. defer s.mu.Unlock()
  195. if s.getErr != nil {
  196. return "", s.getErr
  197. }
  198. value, ok := s.values[jobID]
  199. if !ok {
  200. return "", errors.New("missing")
  201. }
  202. return value, nil
  203. }
  204. func (s *fakePayloadStore) Delete(_ context.Context, jobID int64) error {
  205. s.mu.Lock()
  206. defer s.mu.Unlock()
  207. if s.trace != nil {
  208. *s.trace = append(*s.trace, "payload_delete")
  209. }
  210. s.deleted = append(s.deleted, jobID)
  211. delete(s.values, jobID)
  212. return s.deleteErr
  213. }
  214. func (s *fakePayloadStore) Ping(context.Context) error { return s.pingErr }
  215. func asyncConfig() ActiveConfig {
  216. return ActiveConfig{
  217. RiskControlEnabled: true, Enabled: true, BlockingEnabled: false, Strategy: "priority",
  218. WorkerCount: 1, QueueCapacity: 8, Scanners: []string{"pii"}, AllGroups: true, ConfigVersion: 7,
  219. Endpoints: []ActiveEndpoint{{ID: "guard", Enabled: true, TimeoutMS: 1000, InputLimit: 3}},
  220. }
  221. }
  222. func asyncRequest() Request {
  223. return Request{RequestID: "request-async", Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"user","content":"payload canary text"}]}`)}
  224. }
  225. func TestEnqueuerStagingPayloadPublishProtocolAndFailureCleanup(t *testing.T) {
  226. t.Run("success", func(t *testing.T) {
  227. trace := []string{}
  228. repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 41}}
  229. payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
  230. enqueuer := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload)
  231. require.NoError(t, enqueuer.Enqueue(context.Background(), asyncRequest()))
  232. require.Equal(t, []string{"create_staging", "payload_set", "publish_queued"}, trace)
  233. require.Empty(t, repo.createdSnapshot.ScanText)
  234. require.Equal(t, "payload canary text", payload.values[41])
  235. require.Equal(t, DefaultPayloadTTL, payload.setTTL)
  236. })
  237. t.Run("queue admission failures never touch payload", func(t *testing.T) {
  238. for _, createErr := range []error{ErrQueueFull, ErrQueueAdmissionBusy, errors.New("database down")} {
  239. trace := []string{}
  240. repo := &fakeJobRepository{trace: &trace, createErr: createErr}
  241. payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
  242. err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
  243. require.ErrorIs(t, err, createErr)
  244. require.Equal(t, []string{"create_staging"}, trace)
  245. }
  246. })
  247. t.Run("payload failure marks staging failed", func(t *testing.T) {
  248. trace := []string{}
  249. repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 42}}
  250. payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}, setErr: errors.New("redis down")}
  251. err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
  252. require.Error(t, err)
  253. require.Equal(t, []string{"create_staging", "payload_set", "mark_staging_failed"}, trace)
  254. require.Equal(t, "payload_store_failed", repo.markedCode)
  255. })
  256. t.Run("publish failure removes payload and marks staging failed", func(t *testing.T) {
  257. trace := []string{}
  258. repo := &fakeJobRepository{trace: &trace, createJob: &Job{ID: 43}, publishErr: errors.New("publish down")}
  259. payload := &fakePayloadStore{trace: &trace, values: map[int64]string{}}
  260. err := NewEnqueuer(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload).Enqueue(context.Background(), asyncRequest())
  261. require.Error(t, err)
  262. require.Equal(t, []string{"create_staging", "payload_set", "publish_queued", "payload_delete", "mark_staging_failed"}, trace)
  263. require.Equal(t, "queue_publish_failed", repo.markedCode)
  264. require.NotContains(t, payload.values, int64(43))
  265. })
  266. }
  267. func TestEnqueuerSkipsOffOutOfScopeAndNoText(t *testing.T) {
  268. tests := []struct {
  269. name string
  270. cfg ActiveConfig
  271. req Request
  272. }{
  273. {name: "off", cfg: ActiveConfig{}, req: asyncRequest()},
  274. {name: "out of scope", cfg: func() ActiveConfig {
  275. cfg := asyncConfig()
  276. cfg.AllGroups = false
  277. cfg.GroupIDs = []int64{9}
  278. return cfg
  279. }(), req: asyncRequest()},
  280. {name: "no user text", cfg: asyncConfig(), req: Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"function","content":"not audited"}]}`)}},
  281. }
  282. for _, tt := range tests {
  283. t.Run(tt.name, func(t *testing.T) {
  284. repo := &fakeJobRepository{}
  285. err := NewEnqueuer(&fakeConfigStore{cfg: tt.cfg, active: true}, repo, &fakePayloadStore{}).Enqueue(context.Background(), tt.req)
  286. require.NoError(t, err)
  287. require.Zero(t, repo.createdSnapshot.MessageCount)
  288. })
  289. }
  290. }
  291. func TestEnqueuerRecordsAcceptedDroppedAndSkippedMetrics(t *testing.T) {
  292. t.Run("accepted increments enqueued", func(t *testing.T) {
  293. metrics := NewAtomicMetrics()
  294. repo := &fakeJobRepository{createJob: &Job{ID: 44}}
  295. payload := &fakePayloadStore{values: map[int64]string{}}
  296. require.NoError(t, NewEnqueuer(
  297. &fakeConfigStore{cfg: asyncConfig(), active: true},
  298. repo,
  299. payload,
  300. metrics,
  301. ).Enqueue(context.Background(), asyncRequest()))
  302. require.Equal(t, AuditMetricsSnapshot{Enqueued: 1}, metrics.AuditSnapshot())
  303. })
  304. t.Run("queue full increments dropped", func(t *testing.T) {
  305. metrics := NewAtomicMetrics()
  306. repo := &fakeJobRepository{createErr: ErrQueueFull}
  307. err := NewEnqueuer(
  308. &fakeConfigStore{cfg: asyncConfig(), active: true},
  309. repo,
  310. &fakePayloadStore{},
  311. metrics,
  312. ).Enqueue(context.Background(), asyncRequest())
  313. require.ErrorIs(t, err, ErrQueueFull)
  314. require.Equal(t, AuditMetricsSnapshot{Dropped: 1}, metrics.AuditSnapshot())
  315. })
  316. t.Run("skipped request does not increment dropped", func(t *testing.T) {
  317. metrics := NewAtomicMetrics()
  318. require.NoError(t, NewEnqueuer(
  319. &fakeConfigStore{cfg: ActiveConfig{}, active: true},
  320. &fakeJobRepository{},
  321. &fakePayloadStore{},
  322. metrics,
  323. ).Enqueue(context.Background(), asyncRequest()))
  324. require.Equal(t, AuditMetricsSnapshot{}, metrics.AuditSnapshot())
  325. })
  326. }
  327. func workerJob(attempts, maxAttempts int) *Job {
  328. return &Job{ID: 51, ClaimVersion: 3, Attempts: attempts, MaxAttempts: maxAttempts, ConfigVersion: 7,
  329. Snapshot: PromptSnapshot{RequestID: "worker-request", PromptLength: 6, RedactedPreview: "red***"}}
  330. }
  331. func TestWorkerCompletesPassWithoutEventRefreshesEveryChunkAndDeletesPayload(t *testing.T) {
  332. repo := &fakeJobRepository{}
  333. payload := &fakePayloadStore{values: map[int64]string{51: "abcdef"}}
  334. scannerCalls := 0
  335. scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
  336. scannerCalls++
  337. return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", Categories: []string{}, MatchedScanners: []string{}, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}, GuardEndpointID: endpoint.ID}, nil
  338. })
  339. metrics := NewAtomicMetrics()
  340. runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, scanner, metrics)
  341. runner.clock = fixedClock{now: time.Unix(100, 0).UTC()}
  342. require.NoError(t, runner.processJob(context.Background(), 0, asyncConfig(), workerJob(1, 3)))
  343. require.Equal(t, 2, scannerCalls)
  344. require.Equal(t, 2, repo.refreshes)
  345. require.NotNil(t, repo.completedResult)
  346. require.Equal(t, EventPass, repo.completedResult.Decision)
  347. require.False(t, repo.completedStore)
  348. require.Equal(t, []int64{51}, payload.deleted)
  349. require.Equal(t, int64(1), metrics.Snapshot().Total)
  350. require.Equal(t, int64(1), metrics.Snapshot().Allowed)
  351. }
  352. func TestWorkerRetryBackoffTerminalFailureAndFailover(t *testing.T) {
  353. now := time.Unix(200, 0).UTC()
  354. for _, tt := range []struct {
  355. name string
  356. attempts int
  357. maxAttempts int
  358. err *GuardError
  359. wantRetry bool
  360. wantBackoff time.Duration
  361. }{
  362. {name: "first retry", attempts: 1, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 5 * time.Second},
  363. {name: "second retry", attempts: 2, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 30 * time.Second},
  364. {name: "third retry", attempts: 3, maxAttempts: 4, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}, wantRetry: true, wantBackoff: 2 * time.Minute},
  365. {name: "max attempts", attempts: 3, maxAttempts: 3, err: &GuardError{Code: ErrorCodeUnavailable, Retryable: true}},
  366. {name: "invalid terminal", attempts: 1, maxAttempts: 3, err: &GuardError{Code: ErrorCodeInvalidResponse, Retryable: false}},
  367. } {
  368. t.Run(tt.name, func(t *testing.T) {
  369. repo := &fakeJobRepository{}
  370. payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
  371. metrics := NewAtomicMetrics()
  372. runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
  373. return nil, tt.err
  374. }), metrics)
  375. runner.clock = fixedClock{now: now}
  376. err := runner.processJob(context.Background(), 0, asyncConfig(), workerJob(tt.attempts, tt.maxAttempts))
  377. require.Error(t, err)
  378. if tt.wantRetry {
  379. require.Equal(t, 1, repo.retried)
  380. require.Equal(t, now.Add(tt.wantBackoff), repo.retryAt)
  381. require.Empty(t, payload.deleted)
  382. } else {
  383. require.Equal(t, 1, repo.failed)
  384. require.Equal(t, tt.err.Code, repo.failedCode)
  385. require.Equal(t, []int64{51}, payload.deleted)
  386. }
  387. snapshot := metrics.Snapshot()
  388. require.Equal(t, int64(1), snapshot.Total)
  389. if tt.err.Code == ErrorCodeInvalidResponse {
  390. require.Equal(t, int64(1), snapshot.Invalid)
  391. } else {
  392. require.Equal(t, int64(1), snapshot.Unavailable)
  393. }
  394. })
  395. }
  396. repo := &fakeJobRepository{}
  397. payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
  398. metrics := NewAtomicMetrics()
  399. scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, _ string, _ []string) (*NormalizedResult, error) {
  400. if endpoint.ID == "first" {
  401. return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true}
  402. }
  403. return integrationResult(EventPass), nil
  404. })
  405. cfg := asyncConfig()
  406. cfg.Endpoints = []ActiveEndpoint{{ID: "first", Enabled: true, InputLimit: 10}, {ID: "second", Enabled: true, InputLimit: 10}}
  407. runner := NewRunner(&fakeConfigStore{cfg: cfg, active: true}, repo, payload, scanner, metrics)
  408. require.NoError(t, runner.processJob(context.Background(), 0, cfg, workerJob(1, 3)))
  409. require.Equal(t, int64(1), metrics.Snapshot().Failovers)
  410. }
  411. func TestWorkerPanicLeaseLossAndLifecycleAreContained(t *testing.T) {
  412. t.Run("panic", func(t *testing.T) {
  413. repo := &fakeJobRepository{}
  414. payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
  415. runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
  416. panic("scanner panic canary")
  417. }), NewAtomicMetrics())
  418. require.NotPanics(t, func() { runner.processSafely(context.Background(), 0, asyncConfig(), workerJob(1, 3)) })
  419. _, _, failed, _, _, code, message := runner.Snapshot()
  420. require.Equal(t, int64(1), failed)
  421. require.Equal(t, "worker_panic", code)
  422. require.NotContains(t, message, "canary")
  423. require.Equal(t, 1, repo.failed)
  424. })
  425. t.Run("lease loss", func(t *testing.T) {
  426. repo := &fakeJobRepository{refreshErr: ErrLeaseLost}
  427. payload := &fakePayloadStore{values: map[int64]string{51: "abc"}}
  428. calls := 0
  429. runner := NewRunner(&fakeConfigStore{cfg: asyncConfig(), active: true}, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
  430. calls++
  431. return integrationResult(EventPass), nil
  432. }), NewAtomicMetrics())
  433. require.ErrorIs(t, runner.processJob(context.Background(), 0, asyncConfig(), workerJob(1, 3)), ErrLeaseLost)
  434. require.Zero(t, calls)
  435. require.Zero(t, repo.retried)
  436. require.Zero(t, repo.failed)
  437. })
  438. t.Run("start and shutdown", func(t *testing.T) {
  439. cfg := asyncConfig()
  440. cfg.Enabled = false
  441. configStore := &fakeConfigStore{cfg: cfg, active: true}
  442. repo := &fakeJobRepository{}
  443. payload := &fakePayloadStore{pingErr: errors.New("redis unavailable")}
  444. runner := NewRunner(configStore, repo, payload, PromptScannerFunc(func(context.Context, ActiveEndpoint, string, []string) (*NormalizedResult, error) {
  445. return integrationResult(EventPass), nil
  446. }), NewAtomicMetrics())
  447. require.NoError(t, runner.Start(context.Background()))
  448. require.NoError(t, runner.Start(context.Background()))
  449. _, _, _, _, _, code, _ := runner.Snapshot()
  450. require.Equal(t, "payload_store_unavailable", code)
  451. ctx, cancel := context.WithTimeout(context.Background(), time.Second)
  452. defer cancel()
  453. require.NoError(t, runner.Shutdown(ctx))
  454. require.NoError(t, runner.Shutdown(ctx))
  455. })
  456. t.Run("shutdown timeout is bounded", func(t *testing.T) {
  457. runner := &Runner{}
  458. release := make(chan struct{})
  459. runner.wg.Add(1)
  460. go func() {
  461. defer runner.wg.Done()
  462. <-release
  463. }()
  464. ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
  465. defer cancel()
  466. require.ErrorIs(t, runner.Shutdown(ctx), context.DeadlineExceeded)
  467. close(release)
  468. ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
  469. defer cancel2()
  470. require.NoError(t, runner.Shutdown(ctx2))
  471. })
  472. }
  473. func TestPromptAuditSyntheticAsyncBaseline(t *testing.T) {
  474. const totalRequests = 100
  475. cfg := asyncConfig()
  476. cfg.Endpoints[0].InputLimit = 256
  477. cfg.StorePassEvents = false
  478. repo := &fakeJobRepository{}
  479. payload := &fakePayloadStore{values: make(map[int64]string, totalRequests)}
  480. metrics := NewAtomicMetrics()
  481. knownBenignFindings := 0
  482. knownMaliciousBlocked := 0
  483. scanner := PromptScannerFunc(func(_ context.Context, endpoint ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
  484. switch {
  485. case strings.HasPrefix(chunk, "benign"):
  486. return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, Safety: "Safe", GuardEndpointID: endpoint.ID}, nil
  487. case strings.HasPrefix(chunk, "flag"):
  488. return &NormalizedResult{Decision: EventFlag, RiskLevel: RiskMedium, Action: ActionWarn, Safety: "Controversial", Categories: []string{"politically_sensitive_topics"}, GuardEndpointID: endpoint.ID}, nil
  489. case strings.HasPrefix(chunk, "block"):
  490. knownMaliciousBlocked++
  491. return &NormalizedResult{Decision: EventCritical, RiskLevel: RiskCritical, Action: ActionBlock, Safety: "Unsafe", Categories: []string{"jailbreak"}, GuardEndpointID: endpoint.ID}, nil
  492. case strings.HasPrefix(chunk, "invalid"):
  493. return nil, &GuardError{Code: ErrorCodeInvalidResponse}
  494. default:
  495. return nil, &GuardError{Code: ErrorCodeUnavailable, Retryable: true, Timeout: true}
  496. }
  497. })
  498. runner := NewRunner(&fakeConfigStore{cfg: cfg, active: true}, repo, payload, scanner, metrics)
  499. runner.clock = &advancingClock{now: time.Unix(1_000, 0).UTC(), step: time.Millisecond}
  500. for index := 1; index <= totalRequests; index++ {
  501. text := fmt.Sprintf("benign-%03d", index)
  502. switch {
  503. case index > 90 && index <= 95:
  504. text = fmt.Sprintf("flag-%03d", index)
  505. case index > 95 && index <= 98:
  506. text = fmt.Sprintf("block-%03d", index)
  507. case index == 99:
  508. text = "invalid-099"
  509. case index == 100:
  510. text = "timeout-100"
  511. }
  512. jobID := int64(index)
  513. payload.values[jobID] = text
  514. job := &Job{ID: jobID, ClaimVersion: 1, Attempts: 1, MaxAttempts: 1, ConfigVersion: cfg.ConfigVersion,
  515. Snapshot: PromptSnapshot{RequestID: fmt.Sprintf("baseline-%03d", index), PromptLength: len([]rune(text)), RedactedPreview: "synthetic"}}
  516. err := runner.processJob(context.Background(), 0, cfg, job)
  517. if index <= 98 {
  518. require.NoError(t, err)
  519. } else {
  520. require.Error(t, err)
  521. }
  522. }
  523. snapshot := metrics.Snapshot()
  524. require.Equal(t, int64(totalRequests), snapshot.Total)
  525. require.Equal(t, int64(90), snapshot.Allowed)
  526. require.Equal(t, int64(5), snapshot.Flagged)
  527. require.Equal(t, int64(3), snapshot.Blocked)
  528. require.Equal(t, int64(1), snapshot.Invalid)
  529. require.Equal(t, int64(1), snapshot.Unavailable)
  530. require.Equal(t, int64(1), snapshot.Timeouts)
  531. require.Zero(t, knownBenignFindings)
  532. require.Equal(t, 3, knownMaliciousBlocked)
  533. require.Equal(t, 98, repo.completeCount)
  534. require.Equal(t, 8, repo.eventCount, "store_pass_events=false only grows events for flag/block fixtures")
  535. require.Positive(t, snapshot.LatencyP50MS)
  536. require.LessOrEqual(t, snapshot.LatencyP50MS, snapshot.LatencyP95MS)
  537. require.LessOrEqual(t, snapshot.LatencyP95MS, snapshot.LatencyP99MS)
  538. t.Logf("synthetic async baseline: p50=%dms p95=%dms p99=%dms failure_rate=2%% false_positive_rate=0%% event_growth=8/100", snapshot.LatencyP50MS, snapshot.LatencyP95MS, snapshot.LatencyP99MS)
  539. }
  540. func TestRequestCloneOwnsMutableInputs(t *testing.T) {
  541. groupID := int64(7)
  542. req := Request{Body: []byte("original"), GroupID: &groupID}
  543. clone := req.Clone()
  544. clone.Body[0] = 'X'
  545. *clone.GroupID = 8
  546. require.Equal(t, []byte("original"), req.Body)
  547. require.Equal(t, int64(7), *req.GroupID)
  548. require.False(t, reflect.ValueOf(req.Body).Pointer() == reflect.ValueOf(clone.Body).Pointer())
  549. }