8000 fix(collector)!: prevent dropped-byte counter overflow by dominicnunez · Pull Request #322 · dominicnunez/codex-sdk-go · GitHub
[go: up one dir, main page]

Skip to content
Merged
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
6 changes: 6 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,12 @@ jobs:
CGO_ENABLED: '0'
run: go test -v ./appserver/protocol ./appserver/transport -run 'ErrorCode|RPCError|FuzzInboundRPCError' -count=1

- name: Verify 32-bit collector byte counters
env:
GOARCH: '386'
CGO_ENABLED: '0'
run: go test -v ./appserver -run '^TestCollector(Plan|Command)ByteCounter$' -count=1

test-windows:
runs-on: windows-latest
steps:
Expand Down
44 changes: 44 additions & 0 deletions appserver/collector_byte_counter_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
package appserver_test

import (
"runtime"
"strings"
"testing"

codex "github.com/dominicnunez/codex-sdk-go/appserver"
)

func TestCollectorPlanByteCounter(t *testing.T) {
const chunkBytes = 8 << 20
const retainedBytes = 64 << 10
delta := strings.Repeat("x", chunkBytes-4) + "tail"
collector := codex.NewStreamCollector()
for count := 1; count <= 257; count++ {
collector.Process(&codex.PlanDelta{ThreadID: "thread", TurnID: "turn", ItemID: "plan", Delta: delta}, nil)
if count < 256 {
continue
}
summary := collector.Summary()
want := int64(count)*chunkBytes - retainedBytes
if summary.DroppedLatestPlanTextBytes != want {
t.Errorf("%s: after %d chunks dropped %d bytes, want %d", runtime.GOARCH, count, summary.DroppedLatestPlanTextBytes, want)
}
if summary.LatestPlanText == nil || *summary.LatestPlanText != delta[chunkBytes-retainedBytes:] {
t.Fatal("latest plan did not retain the expected suffix")
}
}
// Each component of the plan identity owns a separate accumulation.
for _, scope := range [][3]string{{"other", "turn", "plan"}, {"other", "next", "plan"}, {"other", "next", "new"}} {
collector.Process(&codex.PlanDelta{ThreadID: scope[0], TurnID: scope[1], ItemID: scope[2], Delta: delta[chunkBytes-retainedBytes-1:]}, nil)
summary := collector.Summary()
if summary.DroppedLatestPlanTextBytes != 1 || summary.LatestPlanItemID == nil || *summary.LatestPlanItemID != scope[2] || summary.LatestPlanText == nil || *summary.LatestPlanText != delta[chunkBytes-retainedBytes:] {
t.Fatalf("plan identity %q did not reset text and byte accounting", scope)
}
}
// A completed plan also replaces text within the same identity.
collector.Process(&codex.ItemCompleted{ThreadID: "other", TurnID: "next", Item: codex.ThreadItemWrapper{Value: &codex.PlanThreadItem{ID: "new", Text: "complete"}}}, nil)
summary := collector.Summary()
if summary.DroppedLatestPlanTextBytes != 0 || summary.LatestPlanText == nil || *summary.LatestPlanText != "complete" {
t.Fatal("completed plan did not reset the accumulated text and count")
}
}
88 changes: 88 additions & 0 deletions appserver/collector_command_long_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
//go:build !race

package appserver_test

import (
"context"
"encoding/json"
"runtime"
"strings"
"testing"
"time"

codex "github.com/dominicnunez/codex-sdk-go/appserver"
)

// This public-path fixture decodes more than 2 GiB of JSON. Race builds
// exercise the long counter workload directly in TestCollectorCommandByteCounterOwner
// and keep the existing smaller public notification/lifecycle cases.
func TestCollectorCommandByteCounter(t *testing.T) {
proc, mock := mockProcess(t)
collector := codex.NewStreamCollector()
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
stream := proc.RunStreamedWithCollector(ctx, codex.RunOptions{Prompt: "synthetic accounting fixture"}, collector)
done := make(chan error, 1)
go func() {
var result error
for _, err := range stream.Events() {
if err != nil {
result = err
}
}
done <- result
}()
waitForRunStreamedReady(t, mock)
const chunkBytes = 8 << 20
const retainedBytes = 64 << 10
// Reuse one frame; the fixture never retains a multi-gigabyte history.
delta := strings.Repeat("x", chunkBytes-4) + "tail"
body := json.RawMessage(`{"threadId":"thread-1","turnId":"turn-1","itemId":"command","delta":"` + delta + `"}`)
for count := 1; count <= 257; count++ {
mock.InjectServerNotification(ctx, codex.Notification{JSONRPC: "2.0", Method: "item/commandExecution/outputDelta", Params: body})
if err := ctx.Err(); err != nil {
t.Fatalf("notification fixture interrupted after %d chunks: %v", count, err)
}
if count < 256 {
continue
}
command := collector.Summary().CommandExecutions["command"]
want := int64(count)*chunkBytes - retainedBytes
if command.DroppedOutputDeltaBytes != want {
t.Errorf("%s: after %d chunks dropped %d bytes, want %d", runtime.GOARCH, count, command.DroppedOutputDeltaBytes, want)
}
if len(command.OutputDeltas) != 1 || command.OutputDeltas[0] != delta[chunkBytes-retainedBytes:] || command.DroppedOutputDeltas != count-1 {
t.Fatal("command did not retain the expected suffix and delta count")
}
}
// Wrong run identities must not alter the high counter or create records.
for _, body := range []json.RawMessage{
json.RawMessage(`{"threadId":"other","turnId":"turn-1","itemId":"command","delta":"wrong"}`),
json.RawMessage(`{"threadId":"thread-1","turnId":"other","itemId":"command","delta":"wrong"}`),
} {
mock.InjectServerNotification(ctx, codex.Notification{JSONRPC: "2.0", Method: "item/commandExecution/outputDelta", Params: body})
}
mock.InjectServerNotification(ctx, codex.Notification{JSONRPC: "2.0", Method: "item/commandExecution/outputDelta", Params: json.RawMessage(`{"threadId":"thread-1","turnId":"turn-1","itemId":"other","delta":"separate"}`)})
mock.InjectServerNotification(ctx, codex.Notification{JSONRPC: "2.0", Method: "item/completed", Params: json.RawMessage(`{"completedAtMs":1,"threadId":"thread-1","turnId":"turn-1","item":{"type":"commandExecution","id":"command","command":"synthetic","commandActions":[],"cwd":"/tmp","status":"completed","aggregatedOutput":"complete"}}`)})
mock.InjectServerNotification(ctx, codex.Notification{JSONRPC: "2.0", Method: "turn/completed", Params: json.RawMessage(`{"threadId":"thread-1","turn":{"id":"turn-1","status":"completed","items":[]}}`)})
select {
case err := <-done:
if err != nil {
t.Fatal(err)
}
case <-ctx.Done():
t.Fatal("stream did not complete")
}
summary := collector.Summary()
command := summary.CommandExecutions["command"]
if !command.Completed || command.AggregatedOutput != "complete" || command.DroppedOutputDeltas != 256 {
t.Fatal("command completion changed the accumulated accounting or final output")
}
if command.DroppedOutputDeltaBytes != int64(257)*chunkBytes-retainedBytes {
t.Fatal("command completion changed the accumulated byte count")
}
other := summary.CommandExecutions["other"]
if len(summary.CommandExecutions) != 2 || other.DroppedOutputDeltaBytes != 0 || other.DroppedOutputDeltas != 0 || len(other.OutputDeltas) != 1 || other.OutputDeltas[0] != "separate" {
t.Fatal("command identities did not retain separate accounting")
}
}
31 changes: 31 additions & 0 deletions appserver/collector_counter_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package appserver

import (
"strings"
"testing"
)

func TestCollectorCommandByteCounterOwner(t *testing.T) {
const chunkBytes = 8 << 20
const retainedBytes = 64 << 10
delta := strings.Repeat("x", chunkBytes-4) + "tail"
collector := NewStreamCollector()
for count := 1; count <= 257; count++ {
collector.processCommandExecutionOutputDelta(CommandExecutionOutputDeltaNotification{ThreadID: "thread", TurnID: "turn", ItemID: "command", Delta: delta})
if count < 256 {
continue
}
command := collector.Summary().CommandExecutions["command"]
if want := int64(count)*chunkBytes - retainedBytes; command.DroppedOutputDeltaBytes != want {
t.Errorf("after %d chunks dropped %d bytes, want %d", count, command.DroppedOutputDeltaBytes, want)
}
if len(command.OutputDeltas) != 1 || command.OutputDeltas[0] != delta[chunkBytes-retainedBytes:] || command.DroppedOutputDeltas != count-1 {
t.Fatal("command did not retain the expected suffix and delta count")
}
}
collector.Process(&ItemCompleted{ThreadID: "thread", TurnID: "turn", Item: ThreadItemWrapper{Value: &CommandExecutionThreadItem{ID: "command", Status: CommandExecutionStatusCompleted}}}, nil)
command := collector.Summary().CommandExecutions["command"]
if !command.Completed || command.DroppedOutputDeltaBytes != int64(257)*chunkBytes-retainedBytes || command.DroppedOutputDeltas != 256 {
t.Fatal("completion changed the accumulated accounting")
}
}
4 changes: 2 additions & 2 deletions appserver/collector_storage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ func TestCollectorTrimmedSuffixOwnsStorage(t *testing.T) {
assertOwnedTrimmedString(t, input, *collector.Summary().LatestPlanText)
collector.processCommandExecutionOutputDelta(CommandExecutionOutputDeltaNotification{ItemID: "command", Delta: input})
command := collector.Summary().CommandExecutions["command"]
if len(command.OutputDeltas) != 1 || command.DroppedOutputDeltaBytes != size-streamCollectorOutputDeltaBytesLimit {
if len(command.OutputDeltas) != 1 || command.DroppedOutputDeltaBytes != int64(size-streamCollectorOutputDeltaBytesLimit) {
t.Fatal("command retention accounting changed")
}
assertOwnedTrimmedString(t, input, command.OutputDeltas[0])
Expand All @@ -55,7 +55,7 @@ func TestCollectorDiscardedHistorySlotsAreCleared(t *testing.T) {
maxBytes, maxEntries = 100, 2
}
retained, total, dropped, droppedSize := appendBoundedStringHistory(backing[:2], 5, "4", 0, 0, maxEntries, maxBytes)
if !reflect.DeepEqual(retained, []string{"22", "4"}) || total != bytes || len(retained) != entries || dropped != droppedEntries || droppedSize != droppedBytes {
if !reflect.DeepEqual(retained, []string{"22", "4"}) || total != bytes || len(retained) != entries || dropped != droppedEntries || droppedSize != int64(droppedBytes) {
t.Fatal("append history semantics changed")
}
if backing[0] != "" {
Expand Down
22 changes: 11 additions & 11 deletions appserver/stream_collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ type CommandExecutionLifecycle struct {
CompletedItem *CommandExecutionThreadItem
OutputDeltas []string
DroppedOutputDeltas int
DroppedOutputDeltaBytes int
DroppedOutputDeltaBytes int64
AggregatedOutput string
}

Expand Down Expand Up @@ -88,7 +88,7 @@ type StreamSummary struct {
LatestPlanText *string
LatestPlanItemID *string
LatestTokenUsage *ThreadTokenUsage
DroppedLatestPlanTextBytes int
DroppedLatestPlanTextBytes int64

NormalizedErrors []NormalizedStreamError
DroppedNormalizedErrors int
Expand All @@ -111,7 +111,7 @@ type StreamCollector struct {
latestPlanItemID *string
latestPlanScope lifecycleID
latestTokenUsage *ThreadTokenUsage
droppedLatestPlanTextBytes int
droppedLatestPlanTextBytes int64

normalizedErrors []NormalizedStreamError
droppedNormalizedErrors int
Expand Down Expand Up @@ -316,7 +316,7 @@ func (c *StreamCollector) mergePlanDeltaLocked(p *PlanDelta) {
c.latestPlanItemID = Ptr(p.ItemID)
combined, droppedBytes := appendBoundedStringSuffix(*c.latestPlanText, p.Delta, streamCollectorPlanTextBytesLimit)
c.latestPlanText = Ptr(combined)
c.droppedLatestPlanTextBytes += droppedBytes
c.droppedLatestPlanTextBytes += int64(droppedBytes)
}

func (c *StreamCollector) ingestStartedItemLocked(threadID string, turnID string, item ThreadItem) {
Expand Down Expand Up @@ -571,7 +571,7 @@ func (c *StreamCollector) setLatestPlanTextLocked(id lifecycleID, text string) {
c.latestPlanScope = id
retained, droppedBytes := retainSuffixWithinByteLimit(text, streamCollectorPlanTextBytesLimit)
c.latestPlanText = Ptr(retained)
c.droppedLatestPlanTextBytes = droppedBytes
c.droppedLatestPlanTextBytes = int64(droppedBytes)
}

func cloneMcpToolCallLifecycle(in McpToolCallLifecycle) McpToolCallLifecycle {
Expand Down Expand Up @@ -609,19 +609,19 @@ func appendBoundedHistory[T any](history []T, next T, dropped int, limit int) ([
return history, dropped + 1
}

func appendBoundedStringHistory(history []string, historyBytes int, next string, droppedEntries int, droppedBytes int, maxEntries int, maxBytes int) ([]string, int, int, int) {
func appendBoundedStringHistory(history []string, historyBytes int, next string, droppedEntries int, droppedBytes int64, maxEntries int, maxBytes int) ([]string, int, int, int64) {
if maxEntries <= 0 || maxBytes <= 0 {
return history, 0, droppedEntries + 1, droppedBytes + len(next)
return history, 0, droppedEntries + 1, droppedBytes + int64(len(next))
}

retainedNext, trimmedBytes := retainSuffixWithinByteLimit(next, maxBytes)
droppedBytes += trimmedBytes
droppedBytes += int64(trimmedBytes)

history = append(history, retainedNext)
historyBytes += len(retainedNext)

for len(history) > maxEntries {
droppedBytes += len(history[0])
droppedBytes += int64(len(history[0]))
historyBytes -= len(history[0])
history[0] = ""
history = history[1:]
Expand All @@ -633,11 +633,11 @@ func appendBoundedStringHistory(history []string, historyBytes int, next string,
trimmed, trimmedFromEntry := retainSuffixWithinByteLimit(history[0], maxBytes)
history[0] = trimmed
historyBytes = len(trimmed)
droppedBytes += trimmedFromEntry
droppedBytes += int64(trimmedFromEntry)
break
}

droppedBytes += len(history[0])
droppedBytes += int64(len(history[0]))
historyBytes -= len(history[0])
history[0] = ""
history = history[1:]
Expand Down
6 changes: 6 additions & 0 deletions docs/architecture.md
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,12 @@ Compatible bare presentation keys need reservation before fallback selection;
test adversarial spellings, insertion-order independence, accounting and
completion cleanup for each tuple.

Plan-text and command-output dropped-byte totals use `int64` throughout stored
state, accumulation and summaries. They count discarded input over the applicable
lifecycle independently of retained-text budgets, including on 32-bit builds.
Latest-plan replacement resets its total; command completion preserves its total.
Retained lengths and dropped-entry counts remain `int`.

## Initialization and credentials

Authorization query parsing owns syntax validation for pasted full redirect
Expand Down
4 changes: 4 additions & 0 deletions docs/threat-model.md
Original file line number Diff line number Diff line change
Expand Up @@ -649,6 +649,10 @@ lifecycle's fields; presentation keys must not be parsed as authority. Correlati
separates records but does not authenticate peer-provided IDs or content.
Latest plan deltas append only within the same full tuple; changing thread, turn
or item identity replaces that latest text and resets its dropped-byte accounting.
Both plan and command dropped-byte totals use 64-bit integers independently of
retained-text budgets, preventing long streams from wrapping those totals at the
32-bit integer boundary. This accounting does not add a retention or event-count
limit; command completion preserves the accumulated dropped-byte total.

For an outbound request admitted to the serialized write queue, accepting its
correlated result and abandoning it on cancellation or transport failure share
Expand Down
Loading
0