Files
letta-server/tests/managers/test_file_manager.py
Kian Jones f5c4ab50f4 chore: add ty + pre-commit hook and repeal even more ruff rules (#9504)
* auto fixes

* auto fix pt2 and transitive deps and undefined var checking locals()

* manual fixes (ignored or letta-code fixed)

* fix circular import

* remove all ignores, add FastAPI rules and Ruff rules

* add ty and precommit

* ruff stuff

* ty check fixes

* ty check fixes pt 2

* error on invalid
2026-02-24 10:55:11 -08:00

1178 lines
45 KiB
Python

import asyncio
import time
import pytest
# Import shared fixtures and constants from conftest
from conftest import (
CREATE_DELAY_SQLITE,
USING_SQLITE,
)
from letta.schemas.file import FileMetadata as PydanticFileMetadata
# ======================================================================================================================
# FileAgent Tests
# ======================================================================================================================
@pytest.mark.asyncio
async def test_attach_creates_association(server, default_user, sarah_agent, default_file):
assoc, _closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
visible_content="hello",
max_files_open=sarah_agent.max_files_open,
)
assert assoc.file_id == default_file.id
assert assoc.is_open is True
assert assoc.visible_content == "hello"
sarah_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
file_blocks = sarah_agent.memory.file_blocks
assert len(file_blocks) == 1
assert file_blocks[0].value == assoc.visible_content
assert file_blocks[0].label == default_file.file_name
async def test_attach_is_idempotent(server, default_user, sarah_agent, default_file):
a1, _closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
visible_content="first",
max_files_open=sarah_agent.max_files_open,
)
# second attach with different params
a2, _closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
is_open=False,
visible_content="second",
max_files_open=sarah_agent.max_files_open,
)
assert a1.id == a2.id
assert a2.is_open is False
assert a2.visible_content == "second"
sarah_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
file_blocks = sarah_agent.memory.file_blocks
assert len(file_blocks) == 1
assert file_blocks[0].value == "" # not open
assert file_blocks[0].label == default_file.file_name
async def test_update_file_agent(server, file_attachment, default_user):
updated = await server.file_agent_manager.update_file_agent_by_id(
agent_id=file_attachment.agent_id,
file_id=file_attachment.file_id,
actor=default_user,
is_open=False,
visible_content="updated",
)
assert updated.is_open is False
assert updated.visible_content == "updated"
async def test_update_file_agent_by_file_name(server, file_attachment, default_user):
updated = await server.file_agent_manager.update_file_agent_by_name(
agent_id=file_attachment.agent_id,
file_name=file_attachment.file_name,
actor=default_user,
is_open=False,
visible_content="updated",
)
assert updated.is_open is False
assert updated.visible_content == "updated"
assert updated.start_line is None # start_line should default to None
assert updated.end_line is None # end_line should default to None
@pytest.mark.asyncio
async def test_file_agent_line_tracking(server, default_user, sarah_agent, default_source):
"""Test that line information is captured when opening files with line ranges"""
from letta.schemas.file import FileMetadata as PydanticFileMetadata
# Create a test file with multiple lines
test_content = "line 1\nline 2\nline 3\nline 4\nline 5"
file_metadata = PydanticFileMetadata(
file_name="test_lines.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=test_content)
# Test opening with line range using enforce_max_open_files_and_open
_closed_files, _was_already_open, previous_ranges = await server.file_agent_manager.enforce_max_open_files_and_open(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content="2: line 2\n3: line 3",
max_files_open=sarah_agent.max_files_open,
start_line=2, # 1-indexed
end_line=4, # exclusive
)
# Retrieve and verify line tracking
retrieved = await server.file_agent_manager.get_file_agent_by_id(
agent_id=sarah_agent.id,
file_id=file.id,
actor=default_user,
)
assert retrieved.start_line == 2
assert retrieved.end_line == 4
assert previous_ranges == {} # No previous range since it wasn't open before
# Test opening without line range - should clear line info and capture previous range
_closed_files, _was_already_open, previous_ranges = await server.file_agent_manager.enforce_max_open_files_and_open(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content="full file content",
max_files_open=sarah_agent.max_files_open,
start_line=None,
end_line=None,
)
# Retrieve and verify line info is cleared
retrieved = await server.file_agent_manager.get_file_agent_by_id(
agent_id=sarah_agent.id,
file_id=file.id,
actor=default_user,
)
assert retrieved.start_line is None
assert retrieved.end_line is None
assert previous_ranges == {file.file_name: (2, 4)} # Should capture the previous range
async def test_mark_access(server, file_attachment, default_user):
old_ts = file_attachment.last_accessed_at
if USING_SQLITE:
time.sleep(CREATE_DELAY_SQLITE)
else:
await asyncio.sleep(0.01)
await server.file_agent_manager.mark_access(
agent_id=file_attachment.agent_id,
file_id=file_attachment.file_id,
actor=default_user,
)
refreshed = await server.file_agent_manager.get_file_agent_by_id(
agent_id=file_attachment.agent_id,
file_id=file_attachment.file_id,
actor=default_user,
)
assert refreshed.last_accessed_at > old_ts
async def test_list_files_and_agents(
server,
default_user,
sarah_agent,
charles_agent,
default_file,
another_file,
):
# default_file ↔ charles (open)
await server.file_agent_manager.attach_file(
agent_id=charles_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
max_files_open=charles_agent.max_files_open,
)
# default_file ↔ sarah (open)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# another_file ↔ sarah (closed)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=another_file.id,
file_name=another_file.file_name,
source_id=another_file.source_id,
actor=default_user,
is_open=False,
max_files_open=sarah_agent.max_files_open,
)
files_for_sarah = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
assert {f.file_id for f in files_for_sarah} == {default_file.id, another_file.id}
open_only = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert {f.file_id for f in open_only} == {default_file.id}
agents_for_default = await server.file_agent_manager.list_agents_for_file(default_file.id, actor=default_user)
assert {a.agent_id for a in agents_for_default} == {sarah_agent.id, charles_agent.id}
sarah_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
file_blocks = sarah_agent.memory.file_blocks
assert len(file_blocks) == 2
charles_agent = await server.agent_manager.get_agent_by_id_async(agent_id=charles_agent.id, actor=default_user)
file_blocks = charles_agent.memory.file_blocks
assert len(file_blocks) == 1
assert file_blocks[0].value == ""
assert file_blocks[0].label == default_file.file_name
@pytest.mark.asyncio
async def test_list_files_for_agent_paginated_basic(
server,
default_user,
sarah_agent,
default_source,
):
"""Test basic pagination functionality."""
# create 5 files and attach them to sarah
for i in range(5):
file_metadata = PydanticFileMetadata(
file_name=f"paginated_file_{i}.txt",
source_id=default_source.id,
organization_id=default_user.organization_id,
)
file = await server.file_manager.create_file(file_metadata, actor=default_user)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# get first page
page1, cursor1, has_more1 = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
limit=3,
)
assert len(page1) == 3
assert has_more1 is True
assert cursor1 is not None
# get second page using cursor
page2, cursor2, has_more2 = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
cursor=cursor1,
limit=3,
)
assert len(page2) == 2 # only 2 files left (5 total - 3 already fetched)
assert has_more2 is False
assert cursor2 is not None
# verify no overlap between pages
page1_ids = {fa.id for fa in page1}
page2_ids = {fa.id for fa in page2}
assert page1_ids.isdisjoint(page2_ids)
@pytest.mark.asyncio
async def test_list_files_for_agent_paginated_filter_open(
server,
default_user,
sarah_agent,
default_source,
):
"""Test pagination with is_open=True filter."""
# create files: 3 open, 2 closed
for i in range(5):
file_metadata = PydanticFileMetadata(
file_name=f"filter_file_{i}.txt",
source_id=default_source.id,
organization_id=default_user.organization_id,
)
file = await server.file_manager.create_file(file_metadata, actor=default_user)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
is_open=(i < 3), # first 3 are open
max_files_open=sarah_agent.max_files_open,
)
# get only open files
open_files, _cursor, has_more = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
is_open=True,
limit=10,
)
assert len(open_files) == 3
assert has_more is False
assert all(fa.is_open for fa in open_files)
@pytest.mark.asyncio
async def test_list_files_for_agent_paginated_filter_closed(
server,
default_user,
sarah_agent,
default_source,
):
"""Test pagination with is_open=False filter."""
# create files: 2 open, 4 closed
for i in range(6):
file_metadata = PydanticFileMetadata(
file_name=f"closed_file_{i}.txt",
source_id=default_source.id,
organization_id=default_user.organization_id,
)
file = await server.file_manager.create_file(file_metadata, actor=default_user)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
is_open=(i < 2), # first 2 are open, rest are closed
max_files_open=sarah_agent.max_files_open,
)
# paginate through closed files
page1, cursor1, has_more1 = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
is_open=False,
limit=2,
)
assert len(page1) == 2
assert has_more1 is True
assert all(not fa.is_open for fa in page1)
# get second page of closed files
page2, _cursor2, has_more2 = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
is_open=False,
cursor=cursor1,
limit=3,
)
assert len(page2) == 2 # only 2 closed files left
assert has_more2 is False
assert all(not fa.is_open for fa in page2)
@pytest.mark.asyncio
async def test_list_files_for_agent_paginated_empty(
server,
default_user,
charles_agent,
):
"""Test pagination with agent that has no files."""
# charles_agent has no files attached in this test
result, cursor, has_more = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=charles_agent.id,
actor=default_user,
limit=10,
)
assert len(result) == 0
assert cursor is None
assert has_more is False
@pytest.mark.asyncio
async def test_list_files_for_agent_paginated_large_limit(
server,
default_user,
sarah_agent,
default_source,
):
"""Test that large limit returns all files without pagination."""
# create 3 files
for i in range(3):
file_metadata = PydanticFileMetadata(
file_name=f"all_files_{i}.txt",
source_id=default_source.id,
organization_id=default_user.organization_id,
)
file = await server.file_manager.create_file(file_metadata, actor=default_user)
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# request with large limit
all_files, cursor, has_more = await server.file_agent_manager.list_files_for_agent_paginated(
agent_id=sarah_agent.id,
actor=default_user,
limit=100,
)
assert len(all_files) == 3
assert has_more is False
assert cursor is not None # cursor is still set to last item
@pytest.mark.asyncio
async def test_detach_file(server, file_attachment, default_user):
await server.file_agent_manager.detach_file(
agent_id=file_attachment.agent_id,
file_id=file_attachment.file_id,
actor=default_user,
)
res = await server.file_agent_manager.get_file_agent_by_id(
agent_id=file_attachment.agent_id,
file_id=file_attachment.file_id,
actor=default_user,
)
assert res is None
async def test_detach_file_bulk(
server,
default_user,
sarah_agent,
charles_agent,
default_source,
):
"""Test bulk deletion of multiple agent-file associations."""
# Create multiple files
files = []
for i in range(3):
file_metadata = PydanticFileMetadata(
file_name=f"test_file_{i}.txt",
source_id=default_source.id,
organization_id=default_user.organization_id,
)
file = await server.file_manager.create_file(file_metadata, actor=default_user)
files.append(file)
# Attach all files to both agents
for file in files:
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
await server.file_agent_manager.attach_file(
agent_id=charles_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
max_files_open=charles_agent.max_files_open,
)
# Verify all files are attached to both agents
sarah_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
charles_files = await server.file_agent_manager.list_files_for_agent(
charles_agent.id, per_file_view_window_char_limit=charles_agent.per_file_view_window_char_limit, actor=default_user
)
assert len(sarah_files) == 3
assert len(charles_files) == 3
# Test 1: Bulk delete specific files from specific agents
agent_file_pairs = [
(sarah_agent.id, files[0].id), # Remove file 0 from sarah
(sarah_agent.id, files[1].id), # Remove file 1 from sarah
(charles_agent.id, files[1].id), # Remove file 1 from charles
]
deleted_count = await server.file_agent_manager.detach_file_bulk(agent_file_pairs=agent_file_pairs, actor=default_user)
assert deleted_count == 3
# Verify the correct files were deleted
sarah_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
charles_files = await server.file_agent_manager.list_files_for_agent(
charles_agent.id, per_file_view_window_char_limit=charles_agent.per_file_view_window_char_limit, actor=default_user
)
# Sarah should only have file 2 left
assert len(sarah_files) == 1
assert sarah_files[0].file_id == files[2].id
# Charles should have files 0 and 2 left
assert len(charles_files) == 2
charles_file_ids = {f.file_id for f in charles_files}
assert charles_file_ids == {files[0].id, files[2].id}
# Test 2: Empty list should return 0 and not fail
deleted_count = await server.file_agent_manager.detach_file_bulk(agent_file_pairs=[], actor=default_user)
assert deleted_count == 0
# Test 3: Attempting to delete already deleted associations should return 0
agent_file_pairs = [
(sarah_agent.id, files[0].id), # Already deleted
(sarah_agent.id, files[1].id), # Already deleted
]
deleted_count = await server.file_agent_manager.detach_file_bulk(agent_file_pairs=agent_file_pairs, actor=default_user)
assert deleted_count == 0
async def test_org_scoping(
server,
default_user,
other_user_different_org,
sarah_agent,
default_file,
):
# attach as default_user
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=default_file.id,
file_name=default_file.file_name,
source_id=default_file.source_id,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# other org should see nothing
files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=other_user_different_org
)
assert files == []
# ======================================================================================================================
# LRU File Management Tests
# ======================================================================================================================
async def test_mark_access_bulk(server, default_user, sarah_agent, default_source):
"""Test that mark_access_bulk updates last_accessed_at for multiple files."""
import time
# Create multiple files and attach them
files = []
for i in range(3):
file_metadata = PydanticFileMetadata(
file_name=f"test_file_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"test content {i}")
files.append(file)
# Attach all files (they'll be open by default)
attached_files = []
for file in files:
file_agent, _closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content=f"content for {file.file_name}",
max_files_open=sarah_agent.max_files_open,
)
attached_files.append(file_agent)
# Get initial timestamps
initial_times = {}
for file_agent in attached_files:
fa = await server.file_agent_manager.get_file_agent_by_id(agent_id=sarah_agent.id, file_id=file_agent.file_id, actor=default_user)
initial_times[fa.file_name] = fa.last_accessed_at
# Wait a moment to ensure timestamp difference
time.sleep(1.1)
# Use mark_access_bulk on subset of files
file_names_to_mark = [files[0].file_name, files[2].file_name]
await server.file_agent_manager.mark_access_bulk(agent_id=sarah_agent.id, file_names=file_names_to_mark, actor=default_user)
# Check that only marked files have updated timestamps
for i, file in enumerate(files):
fa = await server.file_agent_manager.get_file_agent_by_id(agent_id=sarah_agent.id, file_id=file.id, actor=default_user)
if file.file_name in file_names_to_mark:
assert fa.last_accessed_at > initial_times[file.file_name], f"File {file.file_name} should have updated timestamp"
else:
assert fa.last_accessed_at == initial_times[file.file_name], f"File {file.file_name} should not have updated timestamp"
async def test_lru_eviction_on_attach(server, default_user, sarah_agent, default_source):
"""Test that attaching files beyond max_files_open triggers LRU eviction."""
import time
# Use the agent's configured max_files_open
max_files_open = sarah_agent.max_files_open
# Create more files than the limit
files = []
for i in range(max_files_open + 2): # e.g., 7 files for max_files_open=5
file_metadata = PydanticFileMetadata(
file_name=f"lru_test_file_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"test content {i}")
files.append(file)
# Attach files one by one with small delays to ensure different timestamps
attached_files = []
all_closed_files = []
for i, file in enumerate(files):
if i > 0:
time.sleep(0.1) # Small delay to ensure different timestamps
file_agent, closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content=f"content for {file.file_name}",
max_files_open=sarah_agent.max_files_open,
)
attached_files.append(file_agent)
all_closed_files.extend(closed_files)
# Check that we never exceed max_files_open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id,
per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit,
actor=default_user,
is_open_only=True,
)
assert len(open_files) <= max_files_open, f"Should never exceed {max_files_open} open files"
# Should have closed exactly 2 files (e.g., 7 - 5 = 2 for max_files_open=5)
expected_closed_count = len(files) - max_files_open
assert len(all_closed_files) == expected_closed_count, (
f"Should have closed {expected_closed_count} files, but closed: {all_closed_files}"
)
# Check that the oldest files were closed (first N files attached)
expected_closed = [files[i].file_name for i in range(expected_closed_count)]
assert set(all_closed_files) == set(expected_closed), f"Wrong files closed. Expected {expected_closed}, got {all_closed_files}"
# Check that exactly max_files_open files are open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == max_files_open
# Check that the most recently attached files are still open
open_file_names = {f.file_name for f in open_files}
expected_open = {files[i].file_name for i in range(expected_closed_count, len(files))} # last max_files_open files
assert open_file_names == expected_open
async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, default_source):
"""Test that opening a file beyond max_files_open triggers LRU eviction."""
import time
max_files_open = sarah_agent.max_files_open
# Create files equal to the limit
files = []
for i in range(max_files_open + 1): # 6 files for max_files_open=5
file_metadata = PydanticFileMetadata(
file_name=f"open_test_file_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"test content {i}")
files.append(file)
# Attach first max_files_open files
for i in range(max_files_open):
time.sleep(0.1) # Small delay for different timestamps
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=files[i].id,
file_name=files[i].file_name,
source_id=files[i].source_id,
actor=default_user,
visible_content=f"content for {files[i].file_name}",
max_files_open=sarah_agent.max_files_open,
)
# Attach the last file as closed
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=files[-1].id,
file_name=files[-1].file_name,
source_id=files[-1].source_id,
actor=default_user,
is_open=False,
visible_content=f"content for {files[-1].file_name}",
max_files_open=sarah_agent.max_files_open,
)
# All files should be attached but only max_files_open should be open
all_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(all_files) == max_files_open + 1
assert len(open_files) == max_files_open
# Wait a moment
time.sleep(0.1)
# Now "open" the last file using the efficient method
closed_files, _was_already_open, _ = await server.file_agent_manager.enforce_max_open_files_and_open(
agent_id=sarah_agent.id,
file_id=files[-1].id,
file_name=files[-1].file_name,
source_id=files[-1].source_id,
actor=default_user,
visible_content="updated content",
max_files_open=sarah_agent.max_files_open,
)
# Should have closed 1 file (the oldest one)
assert len(closed_files) == 1, f"Should have closed 1 file, got: {closed_files}"
assert closed_files[0] == files[0].file_name, f"Should have closed oldest file {files[0].file_name}"
# Check that exactly max_files_open files are still open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == max_files_open
# Check that the newly opened file is open and the oldest is closed
last_file_agent = await server.file_agent_manager.get_file_agent_by_id(
agent_id=sarah_agent.id, file_id=files[-1].id, actor=default_user
)
first_file_agent = await server.file_agent_manager.get_file_agent_by_id(
agent_id=sarah_agent.id, file_id=files[0].id, actor=default_user
)
assert last_file_agent.is_open is True, "Last file should be open"
assert first_file_agent.is_open is False, "First file should be closed"
async def test_lru_no_eviction_when_reopening_same_file(server, default_user, sarah_agent, default_source):
"""Test that reopening an already open file doesn't trigger unnecessary eviction."""
import time
max_files_open = sarah_agent.max_files_open
# Create files equal to the limit
files = []
for i in range(max_files_open):
file_metadata = PydanticFileMetadata(
file_name=f"reopen_test_file_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"test content {i}")
files.append(file)
# Attach all files (they'll be open)
for i, file in enumerate(files):
time.sleep(0.1) # Small delay for different timestamps
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content=f"content for {file.file_name}",
max_files_open=sarah_agent.max_files_open,
)
# All files should be open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == max_files_open
initial_open_names = {f.file_name for f in open_files}
# Wait a moment
time.sleep(0.1)
# "Reopen" the last file (which is already open)
closed_files, was_already_open, _ = await server.file_agent_manager.enforce_max_open_files_and_open(
agent_id=sarah_agent.id,
file_id=files[-1].id,
file_name=files[-1].file_name,
source_id=files[-1].source_id,
actor=default_user,
visible_content="updated content",
max_files_open=sarah_agent.max_files_open,
)
# Should not have closed any files since we're within the limit
assert len(closed_files) == 0, f"Should not have closed any files when reopening, got: {closed_files}"
assert was_already_open is True, "File should have been detected as already open"
# All the same files should still be open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == max_files_open
final_open_names = {f.file_name for f in open_files}
assert initial_open_names == final_open_names, "Same files should remain open"
async def test_last_accessed_at_updates_correctly(server, default_user, sarah_agent, default_source):
"""Test that last_accessed_at is updated in the correct scenarios."""
import time
# Create and attach a file
file_metadata = PydanticFileMetadata(
file_name="timestamp_test.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text="test content")
file_agent, _closed_files = await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content="initial content",
max_files_open=sarah_agent.max_files_open,
)
initial_time = file_agent.last_accessed_at
time.sleep(1.1)
# Test update_file_agent_by_id updates timestamp
updated_agent = await server.file_agent_manager.update_file_agent_by_id(
agent_id=sarah_agent.id, file_id=file.id, actor=default_user, visible_content="updated content"
)
assert updated_agent.last_accessed_at > initial_time, "update_file_agent_by_id should update timestamp"
time.sleep(1.1)
prev_time = updated_agent.last_accessed_at
# Test update_file_agent_by_name updates timestamp
updated_agent2 = await server.file_agent_manager.update_file_agent_by_name(
agent_id=sarah_agent.id, file_name=file.file_name, actor=default_user, is_open=False
)
assert updated_agent2.last_accessed_at > prev_time, "update_file_agent_by_name should update timestamp"
time.sleep(1.1)
prev_time = updated_agent2.last_accessed_at
# Test mark_access updates timestamp
await server.file_agent_manager.mark_access(agent_id=sarah_agent.id, file_id=file.id, actor=default_user)
final_agent = await server.file_agent_manager.get_file_agent_by_id(agent_id=sarah_agent.id, file_id=file.id, actor=default_user)
assert final_agent.last_accessed_at > prev_time, "mark_access should update timestamp"
async def test_attach_files_bulk_basic(server, default_user, sarah_agent, default_source):
"""Test basic functionality of attach_files_bulk method."""
# Create multiple files
files = []
for i in range(3):
file_metadata = PydanticFileMetadata(
file_name=f"bulk_test_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"content {i}")
files.append(file)
# Create visible content map
visible_content_map = {f"bulk_test_{i}.txt": f"visible content {i}" for i in range(3)}
# Bulk attach files
closed_files = await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id,
files_metadata=files,
visible_content_map=visible_content_map,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# Should not close any files since we're under the limit
assert closed_files == []
# Verify all files are attached and open
attached_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(attached_files) == 3
attached_file_names = {f.file_name for f in attached_files}
expected_names = {f"bulk_test_{i}.txt" for i in range(3)}
assert attached_file_names == expected_names
# Verify visible content is set correctly
for i, attached_file in enumerate(attached_files):
if attached_file.file_name == f"bulk_test_{i}.txt":
assert attached_file.visible_content == f"visible content {i}"
async def test_attach_files_bulk_deduplication(server, default_user, sarah_agent, default_source):
"""Test that attach_files_bulk properly deduplicates files with same names."""
# Create files with same name (different IDs)
file_metadata_1 = PydanticFileMetadata(
file_name="duplicate_test.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file1 = await server.file_manager.create_file(file_metadata=file_metadata_1, actor=default_user, text="content 1")
file_metadata_2 = PydanticFileMetadata(
file_name="duplicate_test.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file2 = await server.file_manager.create_file(file_metadata=file_metadata_2, actor=default_user, text="content 2")
# Try to attach both files (same name, different IDs)
files_to_attach = [file1, file2]
visible_content_map = {"duplicate_test.txt": "visible content"}
# Bulk attach should deduplicate
await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id,
files_metadata=files_to_attach,
visible_content_map=visible_content_map,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# Should only attach one file (deduplicated)
attached_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
assert len(attached_files) == 1
assert attached_files[0].file_name == "duplicate_test.txt"
async def test_attach_files_bulk_lru_eviction(server, default_user, sarah_agent, default_source):
"""Test that attach_files_bulk properly handles LRU eviction without duplicates."""
import time
max_files_open = sarah_agent.max_files_open
# First, fill up to the max with individual files
existing_files = []
for i in range(max_files_open):
file_metadata = PydanticFileMetadata(
file_name=f"existing_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"existing {i}")
existing_files.append(file)
time.sleep(0.05) # Small delay for different timestamps
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=file.id,
file_name=file.file_name,
source_id=file.source_id,
actor=default_user,
visible_content=f"existing content {i}",
max_files_open=sarah_agent.max_files_open,
)
# Verify we're at the limit
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == max_files_open
# Now bulk attach 3 new files (should trigger LRU eviction)
new_files = []
for i in range(3):
file_metadata = PydanticFileMetadata(
file_name=f"new_bulk_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"new content {i}")
new_files.append(file)
visible_content_map = {f"new_bulk_{i}.txt": f"new visible {i}" for i in range(3)}
# Bulk attach should evict oldest files
closed_files = await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id,
files_metadata=new_files,
visible_content_map=visible_content_map,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# Should have closed exactly 3 files (oldest ones)
assert len(closed_files) == 3
# CRITICAL: Verify no duplicates in closed_files list
assert len(closed_files) == len(set(closed_files)), f"Duplicate file names in closed_files: {closed_files}"
# Verify expected files were closed (oldest 3)
expected_closed = {f"existing_{i}.txt" for i in range(3)}
actual_closed = set(closed_files)
assert actual_closed == expected_closed
# Verify we still have exactly max_files_open files open
open_files_after = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files_after) == max_files_open
# Verify the new files are open
open_file_names = {f.file_name for f in open_files_after}
for i in range(3):
assert f"new_bulk_{i}.txt" in open_file_names
async def test_attach_files_bulk_mixed_existing_new(server, default_user, sarah_agent, default_source):
"""Test bulk attach with mix of existing and new files."""
# Create and attach one file individually first
existing_file_metadata = PydanticFileMetadata(
file_name="existing_file.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
existing_file = await server.file_manager.create_file(file_metadata=existing_file_metadata, actor=default_user, text="existing")
await server.file_agent_manager.attach_file(
agent_id=sarah_agent.id,
file_id=existing_file.id,
file_name=existing_file.file_name,
source_id=existing_file.source_id,
actor=default_user,
visible_content="old content",
is_open=False, # Start as closed
max_files_open=sarah_agent.max_files_open,
)
# Create new files
new_files = []
for i in range(2):
file_metadata = PydanticFileMetadata(
file_name=f"new_file_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"new {i}")
new_files.append(file)
# Bulk attach: existing file + new files
files_to_attach = [existing_file, *new_files]
visible_content_map = {
"existing_file.txt": "updated content",
"new_file_0.txt": "new content 0",
"new_file_1.txt": "new content 1",
}
closed_files = await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id,
files_metadata=files_to_attach,
visible_content_map=visible_content_map,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# Should not close any files
assert closed_files == []
# Verify all files are now open
open_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files) == 3
# Verify existing file was updated
existing_file_agent = await server.file_agent_manager.get_file_agent_by_file_name(
agent_id=sarah_agent.id, file_name="existing_file.txt", actor=default_user
)
assert existing_file_agent.is_open is True
assert existing_file_agent.visible_content == "updated content"
async def test_attach_files_bulk_empty_list(server, default_user, sarah_agent):
"""Test attach_files_bulk with empty file list."""
closed_files = await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id, files_metadata=[], visible_content_map={}, actor=default_user, max_files_open=sarah_agent.max_files_open
)
assert closed_files == []
# Verify no files are attached
attached_files = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
assert len(attached_files) == 0
async def test_attach_files_bulk_oversized_bulk(server, default_user, sarah_agent, default_source):
"""Test bulk attach when trying to attach more files than max_files_open allows."""
max_files_open = sarah_agent.max_files_open
# Create more files than the limit allows
oversized_files = []
for i in range(max_files_open + 3): # 3 more than limit
file_metadata = PydanticFileMetadata(
file_name=f"oversized_{i}.txt",
organization_id=default_user.organization_id,
source_id=default_source.id,
)
file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"oversized {i}")
oversized_files.append(file)
visible_content_map = {f"oversized_{i}.txt": f"oversized visible {i}" for i in range(max_files_open + 3)}
# Bulk attach all files (more than limit)
closed_files = await server.file_agent_manager.attach_files_bulk(
agent_id=sarah_agent.id,
files_metadata=oversized_files,
visible_content_map=visible_content_map,
actor=default_user,
max_files_open=sarah_agent.max_files_open,
)
# Should have closed exactly 3 files (the excess)
assert len(closed_files) == 3
# CRITICAL: Verify no duplicates in closed_files list
assert len(closed_files) == len(set(closed_files)), f"Duplicate file names in closed_files: {closed_files}"
# Should have exactly max_files_open files open
open_files_after = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True
)
assert len(open_files_after) == max_files_open
# All files should be attached (some open, some closed)
all_files_after = await server.file_agent_manager.list_files_for_agent(
sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user
)
assert len(all_files_after) == max_files_open + 3