Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion backend/druks/harnesses/pi.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import contextlib
import json
import re
import shlex
Expand Down Expand Up @@ -127,7 +128,12 @@ async def build_invocation(

def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
try:
events = [json.loads(line) for line in result.stdout.splitlines()]
events = []
# result.stdout keeps only the last STDOUT_BUFFER_LIMIT bytes, so the first
# line can be cut, even inside a UTF-8 character.
for line in result.stdout.splitlines():
with contextlib.suppress(ValueError):
events.append(json.loads(line))
# pi repeats an assistant message on turn_end and agent_end;
# message_end is the one place each appears once.
messages = [
Expand Down
41 changes: 29 additions & 12 deletions backend/tests/test_pi.py
Original file line number Diff line number Diff line change
Expand Up @@ -302,13 +302,15 @@ def test_parse_treats_a_missing_result_file_as_invalid_output(tmp_path: Path) ->
_parse(b'{"type":"agent_settled"}', tmp_path)


@pytest.mark.parametrize(
("stdout", "expected_cost"),
[
(b"not json", None),
def test_parse_skips_the_line_the_buffer_trim_cut(tmp_path: Path) -> None:
(tmp_path / "run-1").mkdir()
(tmp_path / "run-1" / "output.json").write_text('{"answer": "done", "count": 3}')
stdout = b"\n".join(
(
b'\x99 done"}]}}',
_message_end(
"assistant",
stopReason="stop",
usage={
"input": 1,
"output": 1,
Expand All @@ -317,18 +319,33 @@ def test_parse_treats_a_missing_result_file_as_invalid_output(tmp_path: Path) ->
"cost": {"total": 0.25},
},
),
0.25,
),
],
)
def test_parse_treats_a_broken_stream_as_invalid_output(
tmp_path: Path, stdout: bytes, expected_cost: float | None
) -> None:
)
)

output = _parse(stdout, tmp_path)

assert _Contract.model_validate(output) == _Contract(answer="done", count=3)
cost_usd, _ = read_cost(tmp_path / "run-1")
assert cost_usd == 0.25


def test_parse_treats_a_broken_stream_as_invalid_output(tmp_path: Path) -> None:
stdout = _message_end(
"assistant",
usage={
"input": 1,
"output": 1,
"cacheRead": 0,
"cacheWrite": 0,
"cost": {"total": 0.25},
},
)

with pytest.raises(HarnessInvalidOutputError, match="no usable stream"):
_parse(stdout, tmp_path)

cost_usd, _ = read_cost(tmp_path / "run-1")
assert cost_usd == expected_cost
assert cost_usd == 0.25


async def test_a_pasted_key_becomes_the_catalog_entry(client, druks_db) -> None:
Expand Down
Loading