okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffaufrefactor(test): share inspection and stream supportcba5ff8

main
branchesmainall refs →
17.6 KiB405 linesraw
1"""Durable transcript content, checkpoint forks, and retained execution traces."""
2
3import json
4import urllib.error
5import urllib.parse
6
7from harness import exclusive, operation_id
8
9from integration_support import IntegrationScenario
10
11
12class HistoryContentTests(IntegrationScenario):
13    # exclusive: restarts the daemon
14    @exclusive
15    def test_tool_checkpoint_forks_without_executing_its_pending_call(self):
16        for protocol in ("responses", "chat_completions"):
17            with self.subTest(protocol=protocol):
18                app = self.app_for(protocol)
19                unrelated = app.session()
20                self.send(app, unrelated, "continue from branch")
21                session = app.session()
22                message = 'write and inspect a file: "\\\n' + "😀e\u0301漢" * 400
23                self.send(app, session, message)
24                entries = self.history(app, session)
25                preview = self.read(app, f"/sessions/{session}?tail=0")["preview"]
26                self.assertEqual(preview["text"], message[:256])
27                self.assertLessEqual(len(preview["text"]), 256)
28                self.assertLessEqual(len(preview["text"].encode("utf-8")), 1024)
29                self.assertTrue(preview["truncated"])
30                self.assertEqual(
31                    preview["transcript_count"],
32                    len(
33                        {
34                            entry["position"]
35                            for entry in entries
36                            if entry["kind"] != "continuation"
37                        }
38                    ),
39                )
40                summary = next(
41                    item
42                    for item in self.read(app, "/sessions?scope=all")["items"]
43                    if item["id"] == session
44                )
45                self.assertEqual(summary["preview"], preview)
46                self.assertEqual((app.workspace / "example.txt").read_text(), "hello\n")
47                self.assertTrue(
48                    any(
49                        e["kind"] == "assistant" and self.entry_text(e) == "finished"
50                        for e in entries
51                    )
52                )
53                self.assertTrue(
54                    any(
55                        e["kind"] == "thinking"
56                        and self.entry_text(e) == "reasoning about the task"
57                        for e in entries
58                    )
59                )
60                tool = next(e for e in entries if e["kind"] == "tool_result")
61                result = next(
62                    part["value"]
63                    for part in tool["content"]
64                    if part["kind"] == "json" and part["field"] == "result"
65                )
66                trace = next(
67                    part["trace"] for part in tool["content"] if part["kind"] == "trace"
68                )
69                self.assertEqual(json.loads(result)["status"], "ok")
70                self.assertTrue(
71                    any(
72                        a["kind"] == "read" and a["target"].endswith("example.txt")
73                        for a in trace["activities"]
74                    )
75                )
76                self.assertTrue(
77                    any(c["path"].endswith("example.txt") for c in trace["changes"])
78                )
79                checkpoint = next(e for e in entries if e["kind"] == "tool_call")
80                result_page = self.read(
81                    app,
82                    f"/sessions/{session}/history?after={checkpoint['position']}&limit=1",
83                )
84                self.assertFalse(
85                    any(entry["kind"] == "tool_call" for entry in result_page["items"])
86                )
87                result_first = next(
88                    entry
89                    for entry in result_page["items"]
90                    if entry["kind"] == "tool_result"
91                )
92                self.assertEqual(result_first["tool"]["name"], "python")
93                self.assertEqual(
94                    result_first["tool"]["tool_call_id"],
95                    checkpoint["tool"]["tool_call_id"],
96                )
97                self.assertEqual(result_first["turn_id"], checkpoint["turn_id"])
98                branch_id = operation_id()
99                with app.api(
100                    f"/sessions/{branch_id}",
101                    {
102                        "kind": "fork",
103                        "source_session_id": session,
104                        "checkpoint_id": checkpoint["checkpoint_id"],
105                    },
106                    method="PUT",
107                    headers={"If-None-Match": "*"},
108                ) as response:
109                    branch = json.load(response)
110                self.assertEqual(
111                    (branch["provider_profile"], branch["model"]),
112                    (app.profile, "fixture-model"),
113                )
114                self.assertEqual(
115                    branch["creation"]["resolved"],
116                    {
117                        key: branch[key]
118                        for key in ("workspace", "provider_profile", "model", "effort")
119                    },
120                )
121                branch_entries = self.history(app, branch_id)
122                branch_preview = self.read(app, f"/sessions/{branch_id}?tail=0")[
123                    "preview"
124                ]
125                self.assertEqual(
126                    branch_preview["transcript_count"],
127                    len(
128                        {
129                            entry["position"]
130                            for entry in branch_entries
131                            if entry["kind"] != "continuation"
132                        }
133                    ),
134                )
135                self.assertEqual(branch_preview["text"], message[:256])
136                self.assertTrue(branch_preview["truncated"])
137                pending_result = next(
138                    e for e in branch_entries if e["kind"] == "tool_result"
139                )
140                self.assertIn(
141                    "not executed after branch checkpoint",
142                    json.dumps(pending_result["content"]),
143                )
144                self.assertEqual(
145                    [
146                        (e["kind"], e["content"], e["turn_id"])
147                        for e in branch_entries
148                        if e["id"] != pending_result["id"]
149                    ],
150                    [
151                        (e["kind"], e["content"], e["turn_id"])
152                        for e in entries
153                        if e["position"] <= checkpoint["position"]
154                    ],
155                )
156                start = len(self.provider.requests)
157                self.send(app, branch_id, "continue from branch")
158                continued = self.history(app, branch_id)
159                continued_preview = self.read(app, f"/sessions/{branch_id}?tail=0")[
160                    "preview"
161                ]
162                self.assertEqual(
163                    continued_preview["transcript_count"],
164                    len(
165                        {
166                            entry["position"]
167                            for entry in continued
168                            if entry["kind"] != "continuation"
169                        }
170                    ),
171                )
172                self.assertGreater(
173                    len(
174                        {
175                            entry["position"]
176                            for entry in continued
177                            if entry["kind"] != "continuation"
178                        }
179                    ),
180                    len(
181                        {
182                            entry["position"]
183                            for entry in branch_entries
184                            if entry["kind"] != "continuation"
185                        }
186                    ),
187                )
188                self.assertEqual(continued_preview["text"], "continue from branch")
189                self.assertFalse(continued_preview["truncated"])
190                requests = self.records(start)
191                self.assertEqual(len(requests), 1)
192                self.assertIn(
193                    "/chat/completions"
194                    if protocol == "chat_completions"
195                    else "/responses",
196                    requests[0]["path"],
197                )
198                inputs = requests[0]["request"][
199                    "messages" if protocol == "chat_completions" else "input"
200                ]
201                outputs = [
202                    i
203                    for i in inputs
204                    if i.get("role") == "tool"
205                    or i.get("type") == "function_call_output"
206                ]
207                self.assertEqual(
208                    [i.get("content", i.get("output")) for i in outputs],
209                    ["not executed after branch checkpoint"],
210                )
211                self.assertEqual(self.history(app, session), entries)
212
213    # exclusive: restarts the daemon
214    @exclusive
215    def test_unicode_history_pages_preserve_content_and_timestamps(self):
216        for protocol in ("responses", "chat_completions"):
217            with self.subTest(protocol=protocol):
218                app = self.app_for(protocol)
219                session = app.session()
220                self.send(app, session, "write and inspect a file")
221                self.send(app, session, "review the result")
222                timestamps = [
223                    (e["kind"], e["content"], e["created_at"])
224                    for e in self.history(app, session)
225                    if e["kind"] in ("user", "assistant")
226                ]
227                self.assertTrue(timestamps)
228                self.assertTrue(
229                    all(isinstance(stamp, str) for _, _, stamp in timestamps)
230                )
231                self.restart(app)
232                self.assertEqual(
233                    [
234                        (e["kind"], e["content"], e["created_at"])
235                        for e in self.history(app, session)
236                        if e["kind"] in ("user", "assistant")
237                    ],
238                    timestamps,
239                )
240
241                large_inputs = {}
242                for index in range(3):
243                    identity = operation_id()
244                    body = f"continue from branch {index}: " + '😀e\u0301\\"\n' * 75_000
245                    with app.api(
246                        f"/sessions/{session}/inputs/{identity}",
247                        {"kind": "message", "text": body},
248                        method="PUT",
249                    ) as response:
250                        json.load(response)
251                    app.idle(session, timeout=90)
252                    large_inputs[identity] = body
253
254                def traverse(path, direction):
255                    collected, tokens = [], set()
256                    for _ in range(100):
257                        with app.api(path) as response:
258                            encoded = response.read()
259                        self.assertLessEqual(len(encoded), 1_048_576)
260                        page = json.loads(encoded)
261                        self.assertTrue(
262                            page["items"], "history cursor failed to advance"
263                        )
264                        collected = (
265                            page["items"] + collected
266                            if direction == "older"
267                            else collected + page["items"]
268                        )
269                        token = page[direction]
270                        if token is None:
271                            ids = [entry["id"] for entry in collected]
272                            self.assertEqual(len(ids), len(set(ids)))
273                            return collected
274                        self.assertNotIn(token, tokens, "history cursor repeated")
275                        tokens.add(token)
276                        path = f"/sessions/{session}/history?limit=200&next={urllib.parse.quote(token, safe='')}"
277                    self.fail("history pagination did not terminate")
278
279                backwards = traverse(f"/sessions/{session}/history?limit=200", "older")
280                forwards = traverse(
281                    f"/sessions/{session}/history?after=0&limit=200", "newer"
282                )
283                self.assertEqual(
284                    [entry["id"] for entry in forwards],
285                    [entry["id"] for entry in backwards],
286                )
287                for entry in backwards:
288                    if entry["input_id"] not in large_inputs:
289                        continue
290                    references = [
291                        part["reference"]
292                        for part in entry["content"]
293                        if part["kind"] == "reference"
294                    ]
295                    self.assertEqual(len(references), 1)
296                    reference = references[0]
297                    pieces, offset, token = [], 0, None
298                    for _ in range(100):
299                        path = reference["url"]
300                        if token is not None:
301                            path += "?" + urllib.parse.urlencode({"next": token})
302                        content = self.read(app, path)
303                        for part in content["parts"]:
304                            self.assertEqual(part["encoding"], "utf8")
305                            self.assertEqual(part["offset_bytes"], offset)
306                            pieces.append(part["text"])
307                            offset += len(part["text"].encode("utf-8"))
308                        token = content["next"]
309                        if token is None:
310                            self.assertTrue(content["parts"][-1]["complete"])
311                            break
312                    else:
313                        self.fail("full content pagination did not terminate")
314                    self.assertEqual(offset, reference["bytes"])
315                    self.assertEqual(
316                        "".join(pieces), large_inputs.pop(entry["input_id"])
317                    )
318                self.assertEqual(large_inputs, {})
319
320    def test_fork_retains_continuation_and_trace_after_source_deletion(self):
321        for protocol in ("responses", "chat_completions"):
322            with self.subTest(protocol=protocol):
323                app = self.app_for(protocol)
324                session = app.session()
325                self.send(app, session, "write and inspect a file")
326                tool = next(
327                    entry
328                    for entry in self.history(app, session)
329                    if entry["kind"] == "tool_result"
330                )
331                trace = next(
332                    part["trace"] for part in tool["content"] if part["kind"] == "trace"
333                )
334                continuation_id = operation_id()
335                with app.api(
336                    f"/sessions/{session}/inputs/{continuation_id}",
337                    {"kind": "continue"},
338                    method="PUT",
339                ) as response:
340                    json.load(response)
341                app.idle(session, timeout=90)
342                continuation = next(
343                    entry
344                    for entry in self.history(app, session)
345                    if entry["kind"] == "continuation"
346                    and entry["input_id"] == continuation_id
347                )
348                self.send(app, session, "review the result")
349                self.assertEqual(
350                    self.read(app, f"/sessions/{session}")["automatic_name"],
351                    "review the result",
352                )
353                backwards = self.history(app, session)
354                trace_checkpoint = next(
355                    entry
356                    for entry in backwards
357                    if entry["kind"] == "user"
358                    and self.entry_text(entry) == "review the result"
359                )
360                retained_id = operation_id()
361                with app.api(
362                    f"/sessions/{retained_id}",
363                    {
364                        "kind": "fork",
365                        "source_session_id": session,
366                        "checkpoint_id": trace_checkpoint["checkpoint_id"],
367                    },
368                    method="PUT",
369                    headers={"If-None-Match": "*"},
370                ) as response:
371                    json.load(response)
372
373                def retained_trace():
374                    history = self.history(app, retained_id)
375                    marker = next(
376                        entry
377                        for entry in history
378                        if entry["kind"] == "continuation"
379                        and entry["input_id"] == continuation_id
380                    )
381                    self.assertEqual(marker["turn_id"], continuation["turn_id"])
382                    self.assertEqual(marker["content"], continuation["content"])
383                    result = next(
384                        entry
385                        for entry in history
386                        if entry["kind"] == "tool_result"
387                        and entry["tool"]["tool_call_id"]
388                        == tool["tool"]["tool_call_id"]
389                    )
390                    return next(
391                        part["trace"]
392                        for part in result["content"]
393                        if part["kind"] == "trace"
394                    )
395
396                self.assertEqual(retained_trace(), trace)
397                resource = f"/sessions/{session}?view=configuration"
398                with app.api(resource) as response:
399                    json.load(response)
400                    revision = response.headers["ETag"]
401                with app.api(
402                    resource, method="DELETE", headers={"If-Match": revision}
403                ) as response:
404                    self.assertEqual(json.load(response)["state"], "complete")
405                self.assertEqual(retained_trace(), trace)