11import { describe , expect , test } from "bun:test"
22import { MessageV2 } from "../../src/session/message-v2"
3+
4+ function toolOutput ( part : MessageV2 . ToolPart ) : string {
5+ return ( part . state as { output ?: string } ) . output ?? ""
6+ }
37import { MessageID , PartID , SessionID } from "../../src/session/schema"
48import {
59 extractFileOps ,
@@ -467,15 +471,15 @@ describe("maskObservations", () => {
467471
468472 const allParts = messages . flatMap ( ( m ) => m . parts ) . filter ( ( p ) : p is MessageV2 . ToolPart => p . type === "tool" )
469473
470- expect ( allParts [ 0 ] . state . output ) . toContain ( "[output masked" )
471- expect ( allParts [ 0 ] . state . output ) . toContain ( "bash" )
472- expect ( allParts [ 1 ] . state . output ) . toContain ( "[output masked" )
473- expect ( allParts [ 1 ] . state . output ) . toContain ( "read" )
474- expect ( allParts [ 2 ] . state . output ) . toContain ( "[output masked" )
475- expect ( allParts [ 2 ] . state . output ) . toContain ( "write" )
474+ expect ( toolOutput ( allParts [ 0 ] ) ) . toContain ( "[output masked" )
475+ expect ( toolOutput ( allParts [ 0 ] ) ) . toContain ( "bash" )
476+ expect ( toolOutput ( allParts [ 1 ] ) ) . toContain ( "[output masked" )
477+ expect ( toolOutput ( allParts [ 1 ] ) ) . toContain ( "read" )
478+ expect ( toolOutput ( allParts [ 2 ] ) ) . toContain ( "[output masked" )
479+ expect ( toolOutput ( allParts [ 2 ] ) ) . toContain ( "write" )
476480
477- expect ( allParts [ 3 ] . state . output ) . toBe ( "output4" )
478- expect ( allParts [ 4 ] . state . output ) . toBe ( "output5" )
481+ expect ( toolOutput ( allParts [ 3 ] ) ) . toBe ( "output4" )
482+ expect ( toolOutput ( allParts [ 4 ] ) ) . toBe ( "output5" )
479483 } )
480484
481485 test ( "masked output includes tool name" , ( ) => {
@@ -484,7 +488,7 @@ describe("maskObservations", () => {
484488 maskObservations ( messages , 0 )
485489
486490 const part = messages [ 0 ] . parts [ 0 ] as MessageV2 . ToolPart
487- expect ( part . state . output ) . toContain ( "bash" )
491+ expect ( toolOutput ( part ) ) . toContain ( "bash" )
488492 } )
489493
490494 test ( "masked output includes file path when present" , ( ) => {
@@ -497,8 +501,8 @@ describe("maskObservations", () => {
497501 maskObservations ( messages , 0 )
498502
499503 const part = messages [ 0 ] . parts [ 0 ] as MessageV2 . ToolPart
500- expect ( part . state . output ) . toContain ( "read" )
501- expect ( part . state . output ) . toContain ( "src/index.ts" )
504+ expect ( toolOutput ( part ) ) . toContain ( "read" )
505+ expect ( toolOutput ( part ) ) . toContain ( "src/index.ts" )
502506 } )
503507
504508 test ( "skips parts without output" , ( ) => {
@@ -513,7 +517,7 @@ describe("maskObservations", () => {
513517 expect ( part1 . state . status ) . toBe ( "running" )
514518
515519 const part2 = messages [ 1 ] . parts [ 0 ] as MessageV2 . ToolPart
516- expect ( part2 . state . output ) . toContain ( "[output masked" )
520+ expect ( toolOutput ( part2 ) ) . toContain ( "[output masked" )
517521 } )
518522
519523 test ( "preserves all when preserveRecentCount exceeds total" , ( ) => {
@@ -526,8 +530,8 @@ describe("maskObservations", () => {
526530
527531 const part1 = messages [ 0 ] . parts [ 0 ] as MessageV2 . ToolPart
528532 const part2 = messages [ 1 ] . parts [ 0 ] as MessageV2 . ToolPart
529- expect ( part1 . state . output ) . toBe ( "output1" )
530- expect ( part2 . state . output ) . toBe ( "output2" )
533+ expect ( toolOutput ( part1 ) ) . toBe ( "output1" )
534+ expect ( toolOutput ( part2 ) ) . toBe ( "output2" )
531535 } )
532536
533537 test ( "masks all when preserveRecentCount is zero" , ( ) => {
@@ -540,8 +544,8 @@ describe("maskObservations", () => {
540544
541545 const part1 = messages [ 0 ] . parts [ 0 ] as MessageV2 . ToolPart
542546 const part2 = messages [ 1 ] . parts [ 0 ] as MessageV2 . ToolPart
543- expect ( part1 . state . output ) . toContain ( "[output masked" )
544- expect ( part2 . state . output ) . toContain ( "[output masked" )
547+ expect ( toolOutput ( part1 ) ) . toContain ( "[output masked" )
548+ expect ( toolOutput ( part2 ) ) . toContain ( "[output masked" )
545549 } )
546550} )
547551
0 commit comments