127 lines
4.1 KiB
Go
127 lines
4.1 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"testing"
|
|
|
|
"ollie/backend"
|
|
)
|
|
|
|
// TestTruncation_LargeResultTruncated verifies that a tool result exceeding
|
|
// the 128KB safety limit is truncated and a hint is appended.
|
|
func TestTruncation_LargeResultTruncated(t *testing.T) {
|
|
large := strings.Repeat("x", defaultToolResultMaxBytes+10_000)
|
|
|
|
be := defaultBE()
|
|
be.respond = toolsStream([]backend.ToolCall{
|
|
{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
|
|
})
|
|
|
|
c := newCore(t, be, nil)
|
|
c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
|
|
return large, nil, nil
|
|
}
|
|
|
|
evs := collectEvents(context.Background(), c, "big result")
|
|
|
|
toolEvs := byRole(evs, "tool")
|
|
if len(toolEvs) != 1 {
|
|
t.Fatalf("tool events = %d; want 1", len(toolEvs))
|
|
}
|
|
result := toolEvs[0]
|
|
if !strings.HasPrefix(result, strings.Repeat("x", defaultToolResultMaxBytes)) {
|
|
t.Errorf("result does not start with %d x's", defaultToolResultMaxBytes)
|
|
}
|
|
if !strings.Contains(result, "HARD LIMIT") {
|
|
t.Errorf("truncation hint missing from result: %q", result[:min(len(result), 80)])
|
|
}
|
|
if !strings.Contains(result, "execute_code") {
|
|
t.Errorf("tool name missing from truncation hint: %q", result[:min(len(result), 120)])
|
|
}
|
|
}
|
|
|
|
// TestTruncation_SmallResultNotTruncated verifies that results within the limit
|
|
// pass through unchanged.
|
|
func TestTruncation_SmallResultNotTruncated(t *testing.T) {
|
|
be := defaultBE()
|
|
be.respond = toolsStream([]backend.ToolCall{
|
|
{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
|
|
})
|
|
|
|
c := newCore(t, be, nil)
|
|
c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
|
|
return "short result", nil, nil
|
|
}
|
|
|
|
evs := collectEvents(context.Background(), c, "small result")
|
|
|
|
toolEvs := byRole(evs, "tool")
|
|
if len(toolEvs) != 1 {
|
|
t.Fatalf("tool events = %d; want 1", len(toolEvs))
|
|
}
|
|
if toolEvs[0] != "short result" {
|
|
t.Errorf("result = %q; want %q", toolEvs[0], "short result")
|
|
}
|
|
}
|
|
|
|
// TestTruncation_ErrorTruncated verifies that error results are subject to the
|
|
// same 128KB safety limit (the original bug: errors bypassed all truncation).
|
|
func TestTruncation_ErrorTruncated(t *testing.T) {
|
|
longErr := strings.Repeat("e", defaultToolResultMaxBytes+10_000)
|
|
|
|
be := defaultBE()
|
|
be.respond = toolsStream([]backend.ToolCall{
|
|
{ID: "1", Name: "execute_code", Arguments: json.RawMessage(`{}`)},
|
|
})
|
|
|
|
c := newCore(t, be, nil)
|
|
c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
|
|
return "", nil, &mockErr{longErr}
|
|
}
|
|
|
|
evs := collectEvents(context.Background(), c, "error result")
|
|
|
|
toolEvs := byRole(evs, "tool")
|
|
if len(toolEvs) != 1 {
|
|
t.Fatalf("tool events = %d; want 1", len(toolEvs))
|
|
}
|
|
if !strings.Contains(toolEvs[0], "HARD LIMIT") {
|
|
t.Errorf("error result was NOT truncated: len=%d", len(toolEvs[0]))
|
|
}
|
|
if len(toolEvs[0]) > defaultToolResultMaxBytes+200 {
|
|
t.Errorf("error result too large after truncation: len=%d", len(toolEvs[0]))
|
|
}
|
|
}
|
|
|
|
// TestTruncation_CachedResultAlreadyTruncated verifies that a cache hit on a
|
|
// previously-truncated result returns the truncated form, not the original.
|
|
func TestTruncation_CachedResultAlreadyTruncated(t *testing.T) {
|
|
large := strings.Repeat("z", defaultToolResultMaxBytes+10_000)
|
|
|
|
be := defaultBE()
|
|
be.respond = multiTurnToolsStream([][]backend.ToolCall{
|
|
{{ID: "1", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}},
|
|
{{ID: "2", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}},
|
|
})
|
|
|
|
c := newCore(t, be, nil)
|
|
c.r.runtime.ClassifyTool = func(string) bool { return true } // read-safe → cacheable
|
|
c.r.runtime.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, []backend.ContentBlock, error) {
|
|
return large, nil, nil
|
|
}
|
|
|
|
evs := collectEvents(context.Background(), c, "cached truncated")
|
|
|
|
toolEvs := byRole(evs, "tool")
|
|
if len(toolEvs) != 2 {
|
|
t.Fatalf("tool events = %d; want 2", len(toolEvs))
|
|
}
|
|
for i, ev := range toolEvs {
|
|
if !strings.Contains(ev, "HARD LIMIT") {
|
|
t.Errorf("event[%d] missing truncation hint: %q", i, ev[:min(len(ev), 80)])
|
|
}
|
|
}
|
|
}
|