solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4from __future__ import annotations
5
6import math
7
8import pytest
9
10from solstone.think.providers import local_budget
11from solstone.think.providers.local import ContextBudgetExceeded
12
13
14def _count_chars(text: str) -> int:
15 return len(text)
16
17
18def test_compute_input_budget_reserves_output_and_clamps():
19 assert local_budget.compute_input_budget(123, window=2000) == 2000 - 123 - 256
20 assert local_budget.compute_input_budget(900, window=2000) == 2000 - 500 - 256
21
22
23def test_estimate_tokens_overcounts_four_char_baseline():
24 text = "x" * 12
25
26 assert local_budget.estimate_tokens(text) == 4
27 assert local_budget.estimate_tokens(text) > math.ceil(len(text) / 4)
28
29
30def test_split_entries_round_trips_realistic_markdown():
31 block = (
32 "preamble before the first header\n"
33 "## 2026-06-23 09:00:00 - 09:05:00\n"
34 "### Transcript\n"
35 "hello\n"
36 "### Screen Activity\n"
37 "browser changed\n"
38 "## 2026-06-23 09:05:00 - 09:10:00\n"
39 "### screen summary\n"
40 "later\n"
41 )
42
43 entries = local_budget.split_entries(block)
44
45 assert "".join(entries) == block
46 assert entries[0] == "preamble before the first header\n"
47 assert entries[1].startswith("## 2026-06-23")
48 assert entries[2].startswith("### Transcript")
49 assert entries[3].startswith("### Screen Activity")
50
51
52def test_dedup_runs_before_clip_and_prevents_marker(monkeypatch):
53 monkeypatch.setattr(local_budget, "context_window_tokens", lambda: 1000)
54 entry = "### Screen Activity\n" + ("same screen\n" * 30)
55 block = entry * 3
56
57 fitted, input_budget = local_budget.fit_contents(
58 block,
59 None,
60 50,
61 count=_count_chars,
62 )
63
64 assert fitted == entry
65 assert input_budget is None
66 assert local_budget.TRUNCATION_MARKER not in fitted
67
68
69def test_fit_contents_clips_oldest_entries_to_tail():
70 chunks = [
71 "## 2026-06-23 09:00:00 - 09:05:00\n",
72 "### Transcript\noldest " + ("o" * 4300) + "\n",
73 "### Screen Activity\nmiddle " + ("m" * 4300) + "\n",
74 "## 2026-06-23 09:05:00 - 09:10:00\n",
75 "### Transcript\nrecent " + ("r" * 4300) + "\n",
76 "### Screen Activity\nlatest " + ("l" * 4300) + "\n",
77 ]
78 block = "".join(chunks)
79 talent_prompt = "talent prompt"
80 system_instruction = "system"
81
82 fitted_contents, input_budget = local_budget.fit_contents(
83 [block, talent_prompt],
84 system_instruction,
85 8192 * 6,
86 count=_count_chars,
87 )
88
89 assert isinstance(fitted_contents, list)
90 fitted_block = fitted_contents[0]
91 assert fitted_contents[1] is talent_prompt
92 assert fitted_block.startswith(local_budget.TRUNCATION_MARKER + "\n\n")
93 assert "oldest " not in fitted_block
94 assert "middle " not in fitted_block
95 assert "recent " in fitted_block
96 assert "latest " in fitted_block
97 assert input_budget == {
98 "clipped": True,
99 "dropped_chars": sum(len(chunk) for chunk in chunks[:3]),
100 "dropped_entries": 3,
101 "budget_tokens": 12032,
102 }
103 assert (
104 _count_chars(system_instruction)
105 + _count_chars(talent_prompt)
106 + _count_chars(fitted_block)
107 <= input_budget["budget_tokens"]
108 )
109
110
111def test_fit_contents_leaves_non_overflow_unmarked(monkeypatch):
112 monkeypatch.setattr(local_budget, "context_window_tokens", lambda: 1000)
113 contents = ["## Segment\n### Transcript\nshort\n", "prompt"]
114
115 fitted, input_budget = local_budget.fit_contents(
116 contents,
117 "system",
118 50,
119 count=_count_chars,
120 )
121
122 assert fitted == contents
123 assert input_budget is None
124 assert local_budget.TRUNCATION_MARKER not in fitted[0]
125
126
127def test_fit_contents_raises_when_preserved_content_exceeds_budget(monkeypatch):
128 monkeypatch.setattr(local_budget, "context_window_tokens", lambda: 400)
129
130 with pytest.raises(ContextBudgetExceeded) as exc:
131 local_budget.fit_contents(
132 "## Segment\n### Transcript\nshort\n",
133 "s" * 200,
134 10,
135 count=_count_chars,
136 )
137
138 assert exc.value.reason_code == "context_budget_exceeded"
139
140
141def test_fit_contents_noops_role_dicts_and_empty_lists():
142 role_messages = [{"role": "user", "content": "hello"}]
143 empty: list[str] = []
144
145 fitted_messages, message_budget = local_budget.fit_contents(
146 role_messages,
147 None,
148 10,
149 count=_count_chars,
150 )
151 fitted_empty, empty_budget = local_budget.fit_contents(
152 empty,
153 None,
154 10,
155 count=_count_chars,
156 )
157
158 assert fitted_messages is role_messages
159 assert message_budget is None
160 assert fitted_empty is empty
161 assert empty_budget is None