Skip to content

Commit c389246

Browse files
committed
feat(trace): parse Pi coding sessions
Copy persisted Pi JSONL safely, follow the selected v3 session branch, and map deterministic usage, tools, conversation, and raw evidence. Implements R-006 and R-007, plus the provider-copy portion of R-005 in Spec 005. Refs: #3520 Assisted-by: pi Signed-off-by: Vibhav Bobade <vibhav.bobde@gmail.com>
1 parent 8b05840 commit c389246

11 files changed

Lines changed: 1502 additions & 5 deletions

File tree

‎app/cli/internal/trace/pi/parse.go‎

Lines changed: 585 additions & 0 deletions
Large diffs are not rendered by default.
Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,256 @@
1+
// Copyright 2026 The Chainloop Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package pi
16+
17+
import (
18+
"context"
19+
"encoding/json"
20+
"fmt"
21+
"os"
22+
"path/filepath"
23+
"testing"
24+
25+
"github.com/chainloop-dev/chainloop/app/cli/internal/trace"
26+
"github.com/chainloop-dev/chainloop/app/cli/internal/trace/state"
27+
"github.com/chainloop-dev/chainloop/internal/schemavalidators"
28+
"github.com/chainloop-dev/chainloop/pkg/attestation/crafter/materials/aicodingsession"
29+
"github.com/stretchr/testify/assert"
30+
"github.com/stretchr/testify/require"
31+
)
32+
33+
const (
34+
testSessionID = "pi-session-1"
35+
testRootID = "root0001"
36+
)
37+
38+
func TestParseSessionEvidence(t *testing.T) {
39+
rawDir := t.TempDir()
40+
fixture, err := os.ReadFile(filepath.Join("testdata", "session-v3.jsonl"))
41+
require.NoError(t, err)
42+
writeRawSession(t, rawDir, fixture)
43+
44+
result, err := New().ParseSession(context.Background(), &trace.ParseOpts{
45+
SessionDir: rawDir,
46+
SessionID: testSessionID,
47+
AgentVersion: "0.80.10",
48+
})
49+
require.NoError(t, err)
50+
51+
assert.Equal(t, "pi", result.Data.Agent.Name)
52+
assert.Equal(t, "0.80.10", result.Data.Agent.Version)
53+
assert.Equal(t, testSessionID, result.Data.Session.ID)
54+
assert.Equal(t, "2026-01-02T03:04:00Z", result.Data.Session.StartedAt)
55+
assert.Equal(t, "2026-01-02T03:04:18Z", result.Data.Session.EndedAt)
56+
assert.Equal(t, 18, result.Data.Session.DurationSeconds)
57+
58+
require.NotNil(t, result.Data.Model)
59+
assert.Equal(t, "gemini-2.5-pro", result.Data.Model.Primary)
60+
assert.Equal(t, "google", result.Data.Model.Provider)
61+
assert.Equal(t, []string{"claude-sonnet-4", "gemini-2.5-pro", "gpt-5"}, result.Data.Model.ModelsUsed)
62+
63+
require.NotNil(t, result.Data.Usage)
64+
assert.Equal(t, 300, result.Data.Usage.InputTokens)
65+
assert.Equal(t, 60, result.Data.Usage.OutputTokens)
66+
assert.Equal(t, 395, result.Data.Usage.TotalTokens)
67+
assert.Equal(t, 30, result.Data.Usage.CacheReadInputTokens)
68+
assert.Equal(t, 7, result.Data.Usage.CacheCreationInputTokens)
69+
assert.InDelta(t, 0.3, result.Data.Usage.EstimatedCostUSD, 0.0000001)
70+
71+
require.NotNil(t, result.Data.ToolsUsed)
72+
assert.Equal(t, []aicodingsession.ToolSummary{
73+
{ToolName: "bash", InvocationCount: 1},
74+
{ToolName: "edit", InvocationCount: 2},
75+
{ToolName: "write", InvocationCount: 1},
76+
}, result.Data.ToolsUsed.Summary)
77+
assert.Equal(t, 4, result.Data.ToolsUsed.TotalInvocations)
78+
79+
require.NotNil(t, result.Data.Conversation)
80+
assert.Equal(t, 7, result.Data.Conversation.TotalMessages)
81+
assert.Equal(t, 3, result.Data.Conversation.UserMessages)
82+
assert.Equal(t, 4, result.Data.Conversation.AssistantMessages)
83+
84+
ids := rawEntryIDs(t, result.Data.RawSession["main"])
85+
assert.Equal(t, []string{
86+
"u0000001", "a0000001", "b0000001", "u0000003", "a0000003", "c0000001",
87+
"cm000002", "u0000004", "a0000004", "q0000001", "a0000005",
88+
}, ids)
89+
assert.NotContains(t, ids, "m0000001")
90+
assert.NotContains(t, ids, "m0000002")
91+
assert.NotContains(t, ids, "x0000001")
92+
assert.NotContains(t, ids, "z0000001")
93+
94+
assert.Equal(t, []string{
95+
"Pi entry on line 17 has invalid usage.input; ignored",
96+
"Pi entry on line 17 has invalid usage.output; ignored",
97+
"Pi entry on line 17 has invalid usage.cacheRead; ignored",
98+
"Pi entry on line 17 has invalid usage.totalTokens; ignored",
99+
"Pi entry on line 17 has invalid usage.cost.total; ignored",
100+
}, result.Data.Warnings)
101+
102+
require.NoError(t, validateEvidenceData(result.Data))
103+
104+
actual, err := json.MarshalIndent(result, "", " ")
105+
require.NoError(t, err)
106+
goldenPath := filepath.Join("testdata", "session-v3.golden.json")
107+
if os.Getenv("UPDATE_GOLDEN") == "1" {
108+
require.NoError(t, os.WriteFile(goldenPath, append(actual, '\n'), 0o600))
109+
}
110+
want, err := os.ReadFile(goldenPath)
111+
require.NoError(t, err)
112+
assert.JSONEq(t, string(want), string(actual))
113+
114+
second, err := json.Marshal(result)
115+
require.NoError(t, err)
116+
third, err := json.Marshal(result)
117+
require.NoError(t, err)
118+
assert.Equal(t, second, third, "evidence ordering must be deterministic")
119+
}
120+
121+
func TestParseSessionRejectsInvalidStructure(t *testing.T) {
122+
header := `{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}`
123+
root := `{"type":"message","id":"root0001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"user"}}`
124+
125+
cases := []struct {
126+
name string
127+
content string
128+
want string
129+
}{
130+
{name: "missing header", content: "", want: "missing header"},
131+
{name: "invalid header", content: `{}`, want: "invalid v3 header"},
132+
{name: "unsupported version", content: `{"type":"session","version":2,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}`, want: "invalid v3 header"},
133+
{name: "mismatched session", content: `{"type":"session","version":3,"id":"other","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}`, want: "session ID mismatch"},
134+
{name: "duplicate ID", content: header + "\n" + root + "\n" + `{"type":"label","id":"root0001","parentId":"root0001","timestamp":"2026-01-02T03:04:02Z"}` + "\n", want: "duplicate entry ID"},
135+
{name: "selected orphan", content: header + "\n" + root + "\n" + `{"type":"label","id":"leaf0001","parentId":"missing1","timestamp":"2026-01-02T03:04:02Z"}` + "\n", want: "missing parent"},
136+
{name: "cycle", content: header + "\n" + `{"type":"label","id":"cycle001","parentId":"cycle002","timestamp":"2026-01-02T03:04:01Z"}` + "\n" + `{"type":"label","id":"cycle002","parentId":"cycle001","timestamp":"2026-01-02T03:04:02Z"}` + "\n", want: "cycle"},
137+
{name: "malformed interior", content: header + "\n{" + "\n" + root + "\n", want: "invalid JSON on line 2"},
138+
{name: "malformed final with newline", content: header + "\n{\n", want: "invalid JSON on line 2"},
139+
{name: "missing entry timestamp", content: header + "\n" + `{"type":"label","id":"leaf0001","parentId":null}` + "\n", want: "missing entry timestamp"},
140+
{name: "invalid entry timestamp", content: header + "\n" + `{"type":"label","id":"leaf0001","parentId":null,"timestamp":"not-a-time"}` + "\n", want: "invalid entry timestamp"},
141+
{name: "structurally invalid final fragment", content: header + "\n" + root + "\n" + `{"type":"label"}`, want: "missing entry ID"},
142+
}
143+
144+
for _, tc := range cases {
145+
t.Run(tc.name, func(t *testing.T) {
146+
rawDir := t.TempDir()
147+
writeRawSession(t, rawDir, []byte(tc.content))
148+
_, err := New().ParseSession(context.Background(), &trace.ParseOpts{SessionDir: rawDir, SessionID: testSessionID})
149+
require.Error(t, err)
150+
assert.ErrorContains(t, err, tc.want)
151+
assert.NotContains(t, err.Error(), rawDir, "errors must not expose transcript paths")
152+
})
153+
}
154+
}
155+
156+
func TestParseSessionAllowsAbandonedOrphan(t *testing.T) {
157+
content := `{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}` + "\n" +
158+
`{"type":"future","id":"orphan01","parentId":"missing1","timestamp":"2026-01-02T03:04:01Z"}` + "\n" +
159+
`{"type":"message","id":"root0001","parentId":null,"timestamp":"2026-01-02T03:04:02Z","message":{"role":"user"}}` + "\n"
160+
161+
result, err := parseBytes(t, []byte(content))
162+
require.NoError(t, err)
163+
assert.Equal(t, []string{testRootID}, rawEntryIDs(t, result.Data.RawSession["main"]))
164+
}
165+
166+
func TestParseSessionFinalLineHandling(t *testing.T) {
167+
header := `{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}`
168+
entry := `{"type":"message","id":"root0001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"user"}}`
169+
170+
t.Run("valid unterminated line", func(t *testing.T) {
171+
result, err := parseBytes(t, []byte(header+"\n"+entry))
172+
require.NoError(t, err)
173+
assert.Empty(t, result.Data.Warnings)
174+
assert.Equal(t, []string{testRootID}, rawEntryIDs(t, result.Data.RawSession["main"]))
175+
})
176+
177+
t.Run("truncated final line", func(t *testing.T) {
178+
result, err := parseBytes(t, []byte(header+"\n"+entry+"\n"+`{"type":"message"`))
179+
require.NoError(t, err)
180+
assert.Equal(t, []string{"ignored truncated final Pi session line 3"}, result.Data.Warnings)
181+
assert.Equal(t, []string{testRootID}, rawEntryIDs(t, result.Data.RawSession["main"]))
182+
})
183+
}
184+
185+
func TestParseSessionNonFiniteUsage(t *testing.T) {
186+
content := `{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}
187+
{"type":"message","id":"first001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"assistant","usage":{"input":1e400,"cost":{"total":1e400}}}}
188+
`
189+
190+
result, err := parseBytes(t, []byte(content))
191+
require.NoError(t, err)
192+
assert.Zero(t, result.Data.Usage.InputTokens)
193+
assert.Zero(t, result.Data.Usage.EstimatedCostUSD)
194+
assert.Equal(t, []string{
195+
"Pi entry on line 2 has invalid usage.input; ignored",
196+
"Pi entry on line 2 has invalid usage.cost.total; ignored",
197+
}, result.Data.Warnings)
198+
assert.NoError(t, validateEvidenceData(result.Data), "ignored usage must not make the evidence uncraftable")
199+
}
200+
201+
func TestParseSessionUsageOverflow(t *testing.T) {
202+
maxInt := int(^uint(0) >> 1)
203+
content := fmt.Sprintf(`{"type":"session","version":3,"id":"pi-session-1","timestamp":"2026-01-02T03:04:00Z","cwd":"/work/repo"}
204+
{"type":"message","id":"first001","parentId":null,"timestamp":"2026-01-02T03:04:01Z","message":{"role":"assistant","usage":{"input":%d,"output":1e3,"cost":{"total":1.7e308}}}}
205+
{"type":"message","id":"second01","parentId":"first001","timestamp":"2026-01-02T03:04:02Z","message":{"role":"assistant","usage":{"input":1,"cost":{"total":1.7e308}}}}
206+
`, maxInt)
207+
208+
result, err := parseBytes(t, []byte(content))
209+
require.NoError(t, err)
210+
assert.Equal(t, maxInt, result.Data.Usage.InputTokens)
211+
assert.Equal(t, 1000, result.Data.Usage.OutputTokens)
212+
assert.InDelta(t, 1.7e308, result.Data.Usage.EstimatedCostUSD, 1e292)
213+
assert.Equal(t, []string{
214+
"Pi entry on line 3 has invalid usage.input; ignored",
215+
"Pi entry on line 3 has invalid usage.cost.total; ignored",
216+
}, result.Data.Warnings)
217+
}
218+
219+
func parseBytes(t *testing.T, content []byte) (*aicodingsession.Evidence, error) {
220+
t.Helper()
221+
rawDir := t.TempDir()
222+
writeRawSession(t, rawDir, content)
223+
return New().ParseSession(context.Background(), &trace.ParseOpts{SessionDir: rawDir, SessionID: testSessionID})
224+
}
225+
226+
func writeRawSession(t *testing.T, rawDir string, content []byte) {
227+
t.Helper()
228+
require.NoError(t, os.MkdirAll(rawDir, 0o755))
229+
//nolint:gosec // rawDir is always t.TempDir and the session ID is fixed.
230+
require.NoError(t, os.WriteFile(state.RawSessionPath(rawDir, testSessionID), content, 0o600))
231+
}
232+
233+
func rawEntryIDs(t *testing.T, entries []json.RawMessage) []string {
234+
t.Helper()
235+
ids := make([]string, 0, len(entries))
236+
for _, entry := range entries {
237+
var envelope struct {
238+
ID string `json:"id"`
239+
}
240+
require.NoError(t, json.Unmarshal(entry, &envelope))
241+
ids = append(ids, envelope.ID)
242+
}
243+
return ids
244+
}
245+
246+
func validateEvidenceData(data aicodingsession.Data) error {
247+
raw, err := json.Marshal(data)
248+
if err != nil {
249+
return err
250+
}
251+
var value any
252+
if err := json.Unmarshal(raw, &value); err != nil {
253+
return err
254+
}
255+
return schemavalidators.ValidateAICodingSession(value, schemavalidators.AICodingSessionVersion0_1)
256+
}

0 commit comments

Comments
 (0)