feat: add isPartial detection, filepath/language to context

src/context/index.ts

@@ -5,10 +5,12 @@ import type {
 	ByteRange,
 	Chunk,
 	ChunkContext,
+	ChunkEntityInfo,
 	ChunkOptions,
 	EntityInfo,
 	ExtractedEntity,
 	ImportInfo,
+	Language,
 	ScopeTree,
 } from '../types'
 import { getSiblings, type SiblingOptions } from './siblings'
@@ -68,6 +70,26 @@ export const getScopeForRange = (
 	return scopeChain
 }
 
+/**
+ * Check if an entity is partial (not fully contained) within a byte range
+ *
+ * An entity is partial if it overlaps with the range but is not fully contained.
+ *
+ * @param entity - The entity to check
+ * @param byteRange - The chunk's byte range
+ * @returns Whether the entity is partial
+ */
+const isEntityPartial = (
+	entity: ExtractedEntity,
+	byteRange: ByteRange,
+): boolean => {
+	// Entity is partial if it starts before the range or ends after the range
+	return (
+		entity.byteRange.start < byteRange.start ||
+		entity.byteRange.end > byteRange.end
+	)
+}
+
 /**
  * Get entities within a byte range
  *
@@ -76,7 +98,7 @@ export const getScopeForRange = (
  *
  * @param byteRange - The byte range to search
  * @param scopeTree - The scope tree
- * @returns Entity info array for entities in range
+ * @returns Entity info array for entities in range with isPartial detection
  */
 export const getEntitiesInRange = (
 	byteRange: ByteRange,
@@ -90,12 +112,17 @@ export const getEntitiesInRange = (
 		)
 	})
 
-	// Map to EntityInfo
-	return overlappingEntities.map((entity) => ({
-		name: entity.name,
-		type: entity.type,
-		signature: entity.signature,
-	}))
+	// Map to ChunkEntityInfo with additional fields
+	return overlappingEntities.map(
+		(entity): ChunkEntityInfo => ({
+			name: entity.name,
+			type: entity.type,
+			signature: entity.signature,
+			docstring: entity.docstring,
+			lineRange: entity.lineRange,
+			isPartial: isEntityPartial(entity, byteRange),
+		}),
+	)
 }
 
 /**
@@ -169,23 +196,38 @@ export const getRelevantImports = (
 	return filteredImports.map(mapToImportInfo)
 }
 
+/**
+ * Options for attaching context to a chunk
+ */
+export interface AttachContextOptions {
+	/** The rebuilt text info for the chunk */
+	text: RebuiltText
+	/** The scope tree for the file */
+	scopeTree: ScopeTree
+	/** Chunking options */
+	options: ChunkOptions
+	/** The chunk index */
+	index: number
+	/** Total number of chunks */
+	totalChunks: number
+	/** File path of the source file (optional) */
+	filepath?: string
+	/** Programming language of the source (optional) */
+	language?: Language
+}
+
 /**
  * Attach context information to a chunk
  *
- * @param text - The rebuilt text info for the chunk
- * @param scopeTree - The scope tree for the file
- * @param options - Chunking options
- * @param index - The chunk index
- * @param totalChunks - Total number of chunks
+ * @param opts - Options containing all parameters for context attachment
  * @returns Effect yielding the complete chunk with context
  */
 export const attachContext = (
-	text: RebuiltText,
-	scopeTree: ScopeTree,
-	options: ChunkOptions,
-	index: number,
-	totalChunks: number,
+	opts: AttachContextOptions,
 ): Effect.Effect<Chunk, ContextError> => {
+	const { text, scopeTree, options, index, totalChunks, filepath, language } =
+		opts
+
 	return Effect.try({
 		try: () => {
 			// Determine context mode
@@ -194,6 +236,8 @@ export const attachContext = (
 			// For 'none' mode, return minimal context
 			if (contextMode === 'none') {
 				const context: ChunkContext = {
+					filepath,
+					language,
 					scope: [],
 					entities: [],
 					siblings: [],
@@ -228,6 +272,8 @@ export const attachContext = (
 			const imports = getRelevantImports(entities, scopeTree, filterImports)
 
 			const context: ChunkContext = {
+				filepath,
+				language,
 				scope,
 				entities,
 				siblings,

src/types.ts

@@ -158,6 +158,19 @@ export interface EntityInfo {
 	signature?: string
 }
 
+/**
+ * Extended entity info for entities within a chunk
+ * Includes additional context like docstring, line range, and partial status
+ */
+export interface ChunkEntityInfo extends EntityInfo {
+	/** Documentation comment if present */
+	docstring?: string | null
+	/** Line range in source (0-indexed, inclusive) */
+	lineRange?: LineRange
+	/** Whether this entity spans multiple chunks (is partial) */
+	isPartial?: boolean
+}
+
 /**
  * Information about a sibling entity
  */
@@ -190,10 +203,14 @@ export interface ImportInfo {
  * Context information for a chunk
  */
 export interface ChunkContext {
-	/** Scope information */
+	/** File path of the source file */
+	filepath?: string
+	/** Programming language of the source */
+	language?: Language
+	/** Scope information (scope chain from current to root) */
 	scope: EntityInfo[]
 	/** Entities within this chunk */
-	entities: EntityInfo[]
+	entities: ChunkEntityInfo[]
 	/** Nearby sibling entities */
 	siblings: SiblingInfo[]
 	/** Relevant imports */

test/scope.test.ts

@@ -468,3 +468,155 @@ func subtract(a, b int) int {
 		expect(cls).toBeDefined()
 	})
 })
+
+// ============================================================================
+// Context Attachment Tests
+// ============================================================================
+
+describe('context attachment', () => {
+	test('getEntitiesInRange returns entities with isPartial flag', async () => {
+		const code = `function foo() { return 1 }
+function bar() { return 2 }
+function baz() { return 3 }`
+		const entities = await getEntities(code, 'typescript')
+		const tree = buildScopeTreeFromEntities(entities)
+
+		// Import the function we need to test
+		const { getEntitiesInRange } = await import('../src/context/index')
+
+		// Get entities for a range that fully contains 'bar' but not 'foo' or 'baz'
+		const barEntity = entities.find((e) => e.name === 'bar')
+		if (barEntity) {
+			const entitiesInRange = getEntitiesInRange(barEntity.byteRange, tree)
+
+			// Should find bar
+			const bar = entitiesInRange.find((e) => e.name === 'bar')
+			expect(bar).toBeDefined()
+			// bar should NOT be partial since we're using its exact range
+			expect(bar?.isPartial).toBe(false)
+		}
+	})
+
+	test('getEntitiesInRange marks partial entities correctly', async () => {
+		const code = `class BigClass {
+  method1() { return 1 }
+  method2() { return 2 }
+  method3() { return 3 }
+}`
+		const entities = await getEntities(code, 'typescript')
+		const tree = buildScopeTreeFromEntities(entities)
+
+		const { getEntitiesInRange } = await import('../src/context/index')
+
+		// Get just method2's range - this should be inside BigClass
+		const method2 = entities.find((e) => e.name === 'method2')
+		if (method2) {
+			const entitiesInRange = getEntitiesInRange(method2.byteRange, tree)
+
+			// method2 should not be partial (its full range is included)
+			const m2 = entitiesInRange.find((e) => e.name === 'method2')
+			expect(m2?.isPartial).toBe(false)
+
+			// BigClass should be partial (we only have a slice of it)
+			const cls = entitiesInRange.find((e) => e.name === 'BigClass')
+			if (cls) {
+				expect(cls.isPartial).toBe(true)
+			}
+		}
+	})
+
+	test('getEntitiesInRange includes docstring and lineRange', async () => {
+		const code = `/**
+ * A test function with docs.
+ */
+function documented() {
+  return 1
+}`
+		const entities = await getEntities(code, 'typescript')
+		const tree = buildScopeTreeFromEntities(entities)
+
+		const { getEntitiesInRange } = await import('../src/context/index')
+
+		const fn = entities.find((e) => e.name === 'documented')
+		if (fn) {
+			const entitiesInRange = getEntitiesInRange(fn.byteRange, tree)
+			const docFn = entitiesInRange.find((e) => e.name === 'documented')
+
+			expect(docFn).toBeDefined()
+			expect(docFn?.lineRange).toBeDefined()
+			// Docstring should be present if extracted
+			if (fn.docstring) {
+				expect(docFn?.docstring).toContain('test function')
+			}
+		}
+	})
+
+	test('attachContext includes filepath and language', async () => {
+		const { Effect } = await import('effect')
+		const { attachContext } = await import('../src/context/index')
+
+		const code = `function test() { return 1 }`
+		const entities = await getEntities(code, 'typescript')
+		const tree = buildScopeTreeFromEntities(entities)
+
+		const fn = entities[0]
+		if (fn) {
+			const mockText = {
+				text: code,
+				byteRange: { start: 0, end: code.length },
+				lineRange: { start: 0, end: 0 },
+			}
+
+			const chunk = await Effect.runPromise(
+				attachContext({
+					text: mockText,
+					scopeTree: tree,
+					options: {},
+					index: 0,
+					totalChunks: 1,
+					filepath: 'test.ts',
+					language: 'typescript',
+				}),
+			)
+
+			expect(chunk.context.filepath).toBe('test.ts')
+			expect(chunk.context.language).toBe('typescript')
+		}
+	})
+
+	test('attachContext respects contextMode none', async () => {
+		const { Effect } = await import('effect')
+		const { attachContext } = await import('../src/context/index')
+
+		const code = `function test() { return 1 }`
+		const entities = await getEntities(code, 'typescript')
+		const tree = buildScopeTreeFromEntities(entities)
+
+		const mockText = {
+			text: code,
+			byteRange: { start: 0, end: code.length },
+			lineRange: { start: 0, end: 0 },
+		}
+
+		const chunk = await Effect.runPromise(
+			attachContext({
+				text: mockText,
+				scopeTree: tree,
+				options: { contextMode: 'none' },
+				index: 0,
+				totalChunks: 1,
+				filepath: 'test.ts',
+				language: 'typescript',
+			}),
+		)
+
+		// Even in 'none' mode, filepath and language should be present
+		expect(chunk.context.filepath).toBe('test.ts')
+		expect(chunk.context.language).toBe('typescript')
+		// But scope, entities, siblings, imports should be empty
+		expect(chunk.context.scope).toEqual([])
+		expect(chunk.context.entities).toEqual([])
+		expect(chunk.context.siblings).toEqual([])
+		expect(chunk.context.imports).toEqual([])
+	})
+})