@@ -468,3 +468,155 @@ func subtract(a, b int) int {
468468 expect ( cls ) . toBeDefined ( )
469469 } )
470470} )
471+
472+ // ============================================================================
473+ // Context Attachment Tests
474+ // ============================================================================
475+
476+ describe ( 'context attachment' , ( ) => {
477+ test ( 'getEntitiesInRange returns entities with isPartial flag' , async ( ) => {
478+ const code = `function foo() { return 1 }
479+ function bar() { return 2 }
480+ function baz() { return 3 }`
481+ const entities = await getEntities ( code , 'typescript' )
482+ const tree = buildScopeTreeFromEntities ( entities )
483+
484+ // Import the function we need to test
485+ const { getEntitiesInRange } = await import ( '../src/context/index' )
486+
487+ // Get entities for a range that fully contains 'bar' but not 'foo' or 'baz'
488+ const barEntity = entities . find ( ( e ) => e . name === 'bar' )
489+ if ( barEntity ) {
490+ const entitiesInRange = getEntitiesInRange ( barEntity . byteRange , tree )
491+
492+ // Should find bar
493+ const bar = entitiesInRange . find ( ( e ) => e . name === 'bar' )
494+ expect ( bar ) . toBeDefined ( )
495+ // bar should NOT be partial since we're using its exact range
496+ expect ( bar ?. isPartial ) . toBe ( false )
497+ }
498+ } )
499+
500+ test ( 'getEntitiesInRange marks partial entities correctly' , async ( ) => {
501+ const code = `class BigClass {
502+ method1() { return 1 }
503+ method2() { return 2 }
504+ method3() { return 3 }
505+ }`
506+ const entities = await getEntities ( code , 'typescript' )
507+ const tree = buildScopeTreeFromEntities ( entities )
508+
509+ const { getEntitiesInRange } = await import ( '../src/context/index' )
510+
511+ // Get just method2's range - this should be inside BigClass
512+ const method2 = entities . find ( ( e ) => e . name === 'method2' )
513+ if ( method2 ) {
514+ const entitiesInRange = getEntitiesInRange ( method2 . byteRange , tree )
515+
516+ // method2 should not be partial (its full range is included)
517+ const m2 = entitiesInRange . find ( ( e ) => e . name === 'method2' )
518+ expect ( m2 ?. isPartial ) . toBe ( false )
519+
520+ // BigClass should be partial (we only have a slice of it)
521+ const cls = entitiesInRange . find ( ( e ) => e . name === 'BigClass' )
522+ if ( cls ) {
523+ expect ( cls . isPartial ) . toBe ( true )
524+ }
525+ }
526+ } )
527+
528+ test ( 'getEntitiesInRange includes docstring and lineRange' , async ( ) => {
529+ const code = `/**
530+ * A test function with docs.
531+ */
532+ function documented() {
533+ return 1
534+ }`
535+ const entities = await getEntities ( code , 'typescript' )
536+ const tree = buildScopeTreeFromEntities ( entities )
537+
538+ const { getEntitiesInRange } = await import ( '../src/context/index' )
539+
540+ const fn = entities . find ( ( e ) => e . name === 'documented' )
541+ if ( fn ) {
542+ const entitiesInRange = getEntitiesInRange ( fn . byteRange , tree )
543+ const docFn = entitiesInRange . find ( ( e ) => e . name === 'documented' )
544+
545+ expect ( docFn ) . toBeDefined ( )
546+ expect ( docFn ?. lineRange ) . toBeDefined ( )
547+ // Docstring should be present if extracted
548+ if ( fn . docstring ) {
549+ expect ( docFn ?. docstring ) . toContain ( 'test function' )
550+ }
551+ }
552+ } )
553+
554+ test ( 'attachContext includes filepath and language' , async ( ) => {
555+ const { Effect } = await import ( 'effect' )
556+ const { attachContext } = await import ( '../src/context/index' )
557+
558+ const code = `function test() { return 1 }`
559+ const entities = await getEntities ( code , 'typescript' )
560+ const tree = buildScopeTreeFromEntities ( entities )
561+
562+ const fn = entities [ 0 ]
563+ if ( fn ) {
564+ const mockText = {
565+ text : code ,
566+ byteRange : { start : 0 , end : code . length } ,
567+ lineRange : { start : 0 , end : 0 } ,
568+ }
569+
570+ const chunk = await Effect . runPromise (
571+ attachContext ( {
572+ text : mockText ,
573+ scopeTree : tree ,
574+ options : { } ,
575+ index : 0 ,
576+ totalChunks : 1 ,
577+ filepath : 'test.ts' ,
578+ language : 'typescript' ,
579+ } ) ,
580+ )
581+
582+ expect ( chunk . context . filepath ) . toBe ( 'test.ts' )
583+ expect ( chunk . context . language ) . toBe ( 'typescript' )
584+ }
585+ } )
586+
587+ test ( 'attachContext respects contextMode none' , async ( ) => {
588+ const { Effect } = await import ( 'effect' )
589+ const { attachContext } = await import ( '../src/context/index' )
590+
591+ const code = `function test() { return 1 }`
592+ const entities = await getEntities ( code , 'typescript' )
593+ const tree = buildScopeTreeFromEntities ( entities )
594+
595+ const mockText = {
596+ text : code ,
597+ byteRange : { start : 0 , end : code . length } ,
598+ lineRange : { start : 0 , end : 0 } ,
599+ }
600+
601+ const chunk = await Effect . runPromise (
602+ attachContext ( {
603+ text : mockText ,
604+ scopeTree : tree ,
605+ options : { contextMode : 'none' } ,
606+ index : 0 ,
607+ totalChunks : 1 ,
608+ filepath : 'test.ts' ,
609+ language : 'typescript' ,
610+ } ) ,
611+ )
612+
613+ // Even in 'none' mode, filepath and language should be present
614+ expect ( chunk . context . filepath ) . toBe ( 'test.ts' )
615+ expect ( chunk . context . language ) . toBe ( 'typescript' )
616+ // But scope, entities, siblings, imports should be empty
617+ expect ( chunk . context . scope ) . toEqual ( [ ] )
618+ expect ( chunk . context . entities ) . toEqual ( [ ] )
619+ expect ( chunk . context . siblings ) . toEqual ( [ ] )
620+ expect ( chunk . context . imports ) . toEqual ( [ ] )
621+ } )
622+ } )
0 commit comments