Skip to content

Commit 66b63e6

Browse files
maruelmaruel
committed
Update Codex and Pi provider DTOs
Update data transfer objects for Codex and Pi providers to support new protocol fields and refine type definitions. For Codex: - Update JSONRPCMessage to include emittedAtMs and handle ID as a raw message. - Expand Thread DTO with section metadata, direct input capabilities, and refined optional fields. - Add cacheWriteInputTokens to token usage breakdown. - Enhance MCP server status notifications with thread IDs and failure reasons. - Add SkillsChangedNotification. For Pi: - Unify CommandType and EventType into a single EventType. - Add EventEntryAppended and EntryAppendedEvent to track session history. - Enhance ToolExecResult and AgentMessage with error details, custom types, and display flags. - Update content blocks to support partial arguments and stream indices.
1 parent d3e1d67 commit 66b63e6

5 files changed

Lines changed: 313 additions & 155 deletions

File tree

‎providers/codex/dto.go‎

Lines changed: 45 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -606,18 +606,19 @@ const (
606606
)
607607

608608
// JSONRPCMessage is the JSON-RPC 2.0 envelope for codex app-server messages.
609-
// Notifications have Method set and ID nil. Responses have ID set.
609+
// Notifications have Method set. Responses have a non-empty ID.
610610
type JSONRPCMessage struct {
611-
JSONRPC string `json:"jsonrpc"`
612-
Method Method `json:"method,omitzero"`
613-
ID *json.RawMessage `json:"id,omitzero"`
614-
Params json.RawMessage `json:"params,omitzero"`
615-
Result json.RawMessage `json:"result,omitzero"`
616-
Error *JSONRPCError `json:"error,omitzero"`
611+
JSONRPC string `json:"jsonrpc"`
612+
EmittedAt base.TimeMS `json:"emittedAtMs,omitzero"`
613+
Method Method `json:"method,omitzero"`
614+
ID json.RawMessage `json:"id,omitzero"`
615+
Params json.RawMessage `json:"params,omitzero"`
616+
Result json.RawMessage `json:"result,omitzero"`
617+
Error *JSONRPCError `json:"error,omitzero"`
617618
}
618619

619620
// IsResponse returns true if this is a response (has an ID).
620-
func (m *JSONRPCMessage) IsResponse() bool { return m.ID != nil }
621+
func (m *JSONRPCMessage) IsResponse() bool { return len(m.ID) != 0 }
621622

622623
// JSONRPCError is a JSON-RPC 2.0 error object.
623624
type JSONRPCError struct {
@@ -926,29 +927,32 @@ type ThreadStartedNotification struct {
926927

927928
// Thread describes a thread in thread/started params.
928929
type Thread struct {
929-
ID string `json:"id"`
930-
SessionID string `json:"sessionId,omitzero"`
931-
ForkedFromID *string `json:"forkedFromId,omitzero"`
932-
ParentThreadID *string `json:"parentThreadId,omitzero"`
933-
CLIVersion string `json:"cliVersion,omitzero"`
934-
CreatedAt base.TimeS `json:"createdAt,omitzero"`
935-
CWD string `json:"cwd,omitzero"`
936-
Ephemeral bool `json:"ephemeral,omitzero"`
937-
GitInfo *GitInfo `json:"gitInfo,omitzero"`
938-
ModelProvider string `json:"modelProvider,omitzero"`
939-
Path *string `json:"path,omitzero"`
940-
Preview string `json:"preview,omitzero"`
941-
Source json.RawMessage `json:"source,omitzero"`
942-
ThreadSource ThreadSource `json:"threadSource,omitzero"`
943-
UpdatedAt base.TimeS `json:"updatedAt,omitzero"`
944-
RecencyAt base.TimeS `json:"recencyAt,omitzero"`
945-
Status ThreadStatus `json:"status,omitzero"`
946-
Name string `json:"name,omitzero"`
947-
Extra json.RawMessage `json:"extra,omitzero"`
948-
HistoryMode string `json:"historyMode,omitzero"`
949-
AgentNickname string `json:"agentNickname,omitzero"`
950-
AgentRole string `json:"agentRole,omitzero"`
951-
Turns []Turn `json:"turns,omitzero"`
930+
ID string `json:"id"`
931+
SessionID string `json:"sessionId,omitzero"`
932+
ForkedFromID string `json:"forkedFromId,omitzero"`
933+
ParentThreadID string `json:"parentThreadId,omitzero"`
934+
CLIVersion string `json:"cliVersion,omitzero"`
935+
CreatedAt base.TimeS `json:"createdAt,omitzero"`
936+
CWD string `json:"cwd,omitzero"`
937+
Ephemeral bool `json:"ephemeral,omitzero"`
938+
GitInfo *GitInfo `json:"gitInfo,omitzero"`
939+
ModelProvider string `json:"modelProvider,omitzero"`
940+
Path string `json:"path,omitzero"`
941+
Preview string `json:"preview,omitzero"`
942+
Source json.RawMessage `json:"source,omitzero"`
943+
ThreadSource ThreadSource `json:"threadSource,omitzero"`
944+
UpdatedAt base.TimeS `json:"updatedAt,omitzero"`
945+
RecencyAt base.TimeS `json:"recencyAt,omitzero"`
946+
Status ThreadStatus `json:"status,omitzero"`
947+
Name string `json:"name,omitzero"`
948+
Extra json.RawMessage `json:"extra,omitzero"`
949+
Section json.RawMessage `json:"section,omitzero"`
950+
SectionEnteredAt json.RawMessage `json:"sectionEnteredAt,omitzero"`
951+
CanAcceptDirectInput bool `json:"canAcceptDirectInput,omitzero"`
952+
HistoryMode string `json:"historyMode,omitzero"`
953+
AgentNickname string `json:"agentNickname,omitzero"`
954+
AgentRole string `json:"agentRole,omitzero"`
955+
Turns []Turn `json:"turns,omitzero"`
952956
}
953957

954958
// GitInfo is optional Git metadata captured for a thread.
@@ -1377,6 +1381,7 @@ type TokenUsageBreakdown struct {
13771381
TotalTokens int64 `json:"totalTokens"`
13781382
InputTokens int64 `json:"inputTokens"`
13791383
CachedInputTokens int64 `json:"cachedInputTokens"`
1384+
CacheWriteInputTokens int64 `json:"cacheWriteInputTokens"`
13801385
OutputTokens int64 `json:"outputTokens"`
13811386
ReasoningOutputTokens int64 `json:"reasoningOutputTokens"`
13821387
}
@@ -1723,6 +1728,7 @@ type RateLimitSnapshot struct {
17231728
Secondary *RateLimitWindow `json:"secondary,omitzero"`
17241729
Credits *CreditsSnapshot `json:"credits,omitzero"`
17251730
IndividualLimit *SpendControlLimitSnapshot `json:"individualLimit,omitzero"`
1731+
SpendControlReached bool `json:"spendControlReached,omitzero"`
17261732
PlanType PlanType `json:"planType,omitzero"`
17271733
RateLimitReachedType RateLimitReachedType `json:"rateLimitReachedType,omitzero"`
17281734
}
@@ -1919,7 +1925,12 @@ type McpServerOauthLoginCompletedNotification struct {
19191925

19201926
// McpServerStatusUpdatedNotification holds params for mcpServer/startupStatus/updated.
19211927
type McpServerStatusUpdatedNotification struct {
1922-
Name string `json:"name"`
1923-
Status McpServerStartupState `json:"status"`
1924-
Error *string `json:"error,omitzero"`
1928+
ThreadID string `json:"threadId"`
1929+
Name string `json:"name"`
1930+
Status McpServerStartupState `json:"status"`
1931+
Error string `json:"error,omitzero"`
1932+
FailureReason string `json:"failureReason,omitzero"`
19251933
}
1934+
1935+
// SkillsChangedNotification holds params for skills/changed notifications.
1936+
type SkillsChangedNotification struct{}

‎providers/codex/internal_test.go‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,113 @@ func TestParseOpts(t *testing.T) {
8282
})
8383
}
8484

85+
func TestJSONRPCMessage(t *testing.T) {
86+
t.Run("notification", func(t *testing.T) {
87+
var m JSONRPCMessage
88+
if err := json.Unmarshal([]byte(`{"method":"thread/started","params":{},"emittedAtMs":1787231281472}`), &m); err != nil {
89+
t.Fatal(err)
90+
}
91+
if m.EmittedAt != base.TimeMS(1787231281472) {
92+
t.Errorf("EmittedAt = %v, want 1787231281472", m.EmittedAt)
93+
}
94+
if m.IsResponse() {
95+
t.Error("IsResponse() = true, want false")
96+
}
97+
})
98+
t.Run("response ID", func(t *testing.T) {
99+
for _, tc := range []struct {
100+
name string
101+
data string
102+
want string
103+
ok bool
104+
}{
105+
{name: "omitted", data: `{}`, want: "", ok: false},
106+
{name: "null", data: `{"id":null}`, want: "null", ok: true},
107+
{name: "value", data: `{"id":1}`, want: "1", ok: true},
108+
} {
109+
t.Run(tc.name, func(t *testing.T) {
110+
var m JSONRPCMessage
111+
if err := json.Unmarshal([]byte(tc.data), &m); err != nil {
112+
t.Fatal(err)
113+
}
114+
if string(m.ID) != tc.want || m.IsResponse() != tc.ok {
115+
t.Errorf("message = %#v, want ID %q and IsResponse %t", m, tc.want, tc.ok)
116+
}
117+
})
118+
}
119+
})
120+
}
121+
122+
func TestRecordedNotificationFields(t *testing.T) {
123+
t.Run("thread", func(t *testing.T) {
124+
var notification ThreadStartedNotification
125+
input := `{"thread":{"id":"thread","forkedFromId":null,"parentThreadId":"parent","section":null,"sectionEnteredAt":null,"canAcceptDirectInput":true}}`
126+
if err := json.Unmarshal([]byte(input), &notification); err != nil {
127+
t.Fatal(err)
128+
}
129+
if !notification.Thread.CanAcceptDirectInput || notification.Thread.ForkedFromID != "" || notification.Thread.ParentThreadID != "parent" {
130+
t.Errorf("Thread = %#v, want direct input and value optional IDs", notification.Thread)
131+
}
132+
})
133+
t.Run("token usage", func(t *testing.T) {
134+
var notification ThreadTokenUsageUpdatedNotification
135+
input := `{"threadId":"thread","turnId":"turn","tokenUsage":{"total":{"cacheWriteInputTokens":1},"last":{"cacheWriteInputTokens":2}}}`
136+
if err := json.Unmarshal([]byte(input), &notification); err != nil {
137+
t.Fatal(err)
138+
}
139+
if notification.TokenUsage.Total.CacheWriteInputTokens != 1 || notification.TokenUsage.Last.CacheWriteInputTokens != 2 {
140+
t.Errorf("TokenUsage = %#v, want cache-write token counts", notification.TokenUsage)
141+
}
142+
})
143+
t.Run("MCP startup", func(t *testing.T) {
144+
var notification McpServerStatusUpdatedNotification
145+
input := `{"threadId":"thread","name":"node","status":"starting","error":null,"failureReason":null}`
146+
if err := json.Unmarshal([]byte(input), &notification); err != nil {
147+
t.Fatal(err)
148+
}
149+
if notification.ThreadID != "thread" || notification.Error != "" || notification.FailureReason != "" {
150+
t.Errorf("notification = %#v, want empty optional errors", notification)
151+
}
152+
})
153+
t.Run("rate limit", func(t *testing.T) {
154+
for _, tc := range []struct {
155+
name string
156+
data string
157+
want bool
158+
}{
159+
{name: "omitted", data: `{}`, want: false},
160+
{name: "null", data: `{"spendControlReached":null}`, want: false},
161+
{name: "value", data: `{"spendControlReached":true}`, want: true},
162+
} {
163+
t.Run(tc.name, func(t *testing.T) {
164+
var snapshot RateLimitSnapshot
165+
if err := json.Unmarshal([]byte(tc.data), &snapshot); err != nil {
166+
t.Fatal(err)
167+
}
168+
if snapshot.SpendControlReached != tc.want {
169+
t.Errorf("SpendControlReached = %t, want %t", snapshot.SpendControlReached, tc.want)
170+
}
171+
})
172+
}
173+
})
174+
t.Run("MCP startup value errors", func(t *testing.T) {
175+
var notification McpServerStatusUpdatedNotification
176+
input := `{"threadId":"thread","name":"node","status":"failed","error":"failed","failureReason":"missing binary"}`
177+
if err := json.Unmarshal([]byte(input), &notification); err != nil {
178+
t.Fatal(err)
179+
}
180+
if notification.Error != "failed" || notification.FailureReason != "missing binary" {
181+
t.Errorf("notification = %#v, want value optional errors", notification)
182+
}
183+
})
184+
t.Run("skills changed", func(t *testing.T) {
185+
var notification SkillsChangedNotification
186+
if err := json.Unmarshal([]byte(`{}`), &notification); err != nil {
187+
t.Fatal(err)
188+
}
189+
})
190+
}
191+
85192
func TestNotificationTimeMS(t *testing.T) {
86193
t.Run("item_started", func(t *testing.T) {
87194
const input = `{"item":{"id":"u1","type":"userMessage"},"threadId":"t1","turnId":"turn_1","startedAtMs":1780832660165}`

‎providers/pi/client.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,11 @@ func Scoreboard() scoreboard.Score {
6363
return s
6464
}
6565

66+
// commandDiscriminator is a command without parameters or a correlation ID.
67+
type commandDiscriminator struct {
68+
Type EventType `json:"type"`
69+
}
70+
6671
// cmdExecutor is the production executor backed by exec.Cmd.
6772
type cmdExecutor struct{ bin string }
6873

@@ -215,7 +220,7 @@ func (c *Client) ListModels(ctx context.Context) ([]genai.Model, error) {
215220
}()
216221

217222
sc := newScanner(stdout)
218-
if err := msgutil.WriteNDJSON(stdin, GetModelsCmd{Type: CmdGetModels}); err != nil {
223+
if err := msgutil.WriteNDJSON(stdin, commandDiscriminator{Type: CmdGetModels}); err != nil {
219224
return nil, fmt.Errorf("write get_available_models: %w", err)
220225
}
221226
resp, err := readResponseForCommand(sc, CmdGetModels)
@@ -410,7 +415,7 @@ func msgToPromptParts(msg *genai.Message) (string, []ImageContent, error) {
410415
}
411416

412417
// readResponseForCommand reads lines until a response for the given command is found.
413-
func readResponseForCommand(sc *bufio.Scanner, cmd CommandType) (*Response, error) {
418+
func readResponseForCommand(sc *bufio.Scanner, cmd EventType) (*Response, error) {
414419
for sc.Scan() {
415420
var probe LineProbe
416421
if json.Unmarshal(sc.Bytes(), &probe) != nil {

0 commit comments

Comments
 (0)