context overflow prevention
This commit is contained in:
Luxferre
+205
-1
@@ -268,7 +268,7 @@ func clearBantamEnv(t *testing.T) {
|
||||
"BANTAM_API_KEY", "BANTAM_STREAM", "BANTAM_COLOR", "BANTAM_NO_COLOR",
|
||||
"BANTAM_TIMEOUT", "BANTAM_SHELL_TIMEOUT", "BANTAM_MAX_AL_ITERATIONS",
|
||||
"BANTAM_CONTEXT_WINDOW", "BANTAM_REASONING_EFFORT", "BANTAM_TOOLS_DIR",
|
||||
"BANTAM_SKILLS_DIR",
|
||||
"BANTAM_SKILLS_DIR", "BANTAM_MAX_TOOL_RES",
|
||||
} {
|
||||
t.Setenv(k, "")
|
||||
}
|
||||
@@ -2945,3 +2945,207 @@ func TestLLMContextTokensIncludesReasoning(t *testing.T) {
|
||||
t.Errorf("ctxTokens=%d should equal ContextTokens=%d when prompt_tokens is absent", ctxTokens(u), u.ContextTokens)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ---------- context window overflow prevention (max_tool_res) ----------
|
||||
|
||||
func TestMaxToolResDefault(t *testing.T) {
|
||||
clearBantamEnv(t)
|
||||
cfg := getCfg(filepath.Join(t.TempDir(), "missing.cfg"))
|
||||
if cfg.MaxToolRes != 15000 {
|
||||
t.Errorf("default MaxToolRes = %d, want 15000", cfg.MaxToolRes)
|
||||
}
|
||||
if cfg.Raw["max_tool_res"] != "15000" {
|
||||
t.Errorf("default Raw[max_tool_res] = %q, want 15000", cfg.Raw["max_tool_res"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxToolResFileParse(t *testing.T) {
|
||||
clearBantamEnv(t)
|
||||
p := writeCfg(t, "max_tool_res=4096\n")
|
||||
if got := getCfg(p).MaxToolRes; got != 4096 {
|
||||
t.Errorf("max_tool_res from file = %d, want 4096", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxToolResEnvFallback(t *testing.T) {
|
||||
clearBantamEnv(t)
|
||||
t.Setenv("BANTAM_MAX_TOOL_RES", "2048")
|
||||
if got := getCfg(filepath.Join(t.TempDir(), "missing.cfg")).MaxToolRes; got != 2048 {
|
||||
t.Errorf("max_tool_res from env = %d, want 2048", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The .bantam.cfg file takes precedence over the env var (as for every other
|
||||
// tunable), matching the precedence rule stated for max_tool_res.
|
||||
func TestMaxToolResFileOverridesEnv(t *testing.T) {
|
||||
clearBantamEnv(t)
|
||||
t.Setenv("BANTAM_MAX_TOOL_RES", "1111")
|
||||
p := filepath.Join(t.TempDir(), ".bantam.cfg")
|
||||
if err := os.WriteFile(p, []byte("max_tool_res=2222\n"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := getCfg(p).MaxToolRes; got != 2222 {
|
||||
t.Errorf("file should override env: got %d, want 2222", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxToolResIsInternalKey(t *testing.T) {
|
||||
if !internalKey("max_tool_res") {
|
||||
t.Errorf("internalKey(\"max_tool_res\") = false, want true (must not be forwarded to the API)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolResDir(t *testing.T) {
|
||||
t.Setenv("TMPDIR", "/custom/tmp")
|
||||
if got := toolResDir(); got != filepath.Join("/custom/tmp", "bantam", "toolres") {
|
||||
t.Errorf("toolResDir = %q, want /custom/tmp/bantam/toolres", got)
|
||||
}
|
||||
t.Setenv("TMPDIR", "")
|
||||
if got := toolResDir(); got != filepath.Join("/tmp", "bantam", "toolres") {
|
||||
t.Errorf("toolResDir (no TMPDIR) = %q, want /tmp/bantam/toolres", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Small results pass through untouched and create no files.
|
||||
func TestOffloadToolResultSmallPassthrough(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
var tmps []string
|
||||
res := "hello\n\nexit: 0"
|
||||
if got := offloadToolResult(res, 15000, &tmps); got != res {
|
||||
t.Errorf("small result altered: got %q, want %q", got, res)
|
||||
}
|
||||
if len(tmps) != 0 {
|
||||
t.Errorf("small result should not create temp files, got %v", tmps)
|
||||
}
|
||||
}
|
||||
|
||||
// Large results are spilled to $TMPDIR/bantam/toolres and replaced by an
|
||||
// instruction that names the file and its total length; cleanupTemps then
|
||||
// removes the file.
|
||||
func TestOffloadToolResultLargeSpillsAndCleans(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
t.Setenv("TMPDIR", tmp)
|
||||
var tmps []string
|
||||
big := strings.Repeat("x", 20000)
|
||||
got := offloadToolResult(big, 15000, &tmps)
|
||||
if len(tmps) != 1 {
|
||||
t.Fatalf("expected 1 temp file, got %d (%v)", len(tmps), tmps)
|
||||
}
|
||||
path := tmps[0]
|
||||
wantDir := filepath.Join(tmp, "bantam", "toolres")
|
||||
if filepath.Dir(path) != wantDir {
|
||||
t.Errorf("temp file dir = %q, want %q", filepath.Dir(path), wantDir)
|
||||
}
|
||||
// The full result must be recoverable from the file.
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read spilled file: %v", err)
|
||||
}
|
||||
if string(data) != big {
|
||||
t.Errorf("spilled file content mismatch (len %d vs %d)", len(data), len(big))
|
||||
}
|
||||
// The replacement must tell the model the path and the byte length.
|
||||
if strings.Contains(got, big) {
|
||||
t.Errorf("replacement still contains the full result")
|
||||
}
|
||||
if !strings.Contains(got, path) {
|
||||
t.Errorf("replacement does not reference the file path: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "20000") {
|
||||
t.Errorf("replacement does not state the file length: %q", got)
|
||||
}
|
||||
cleanupTemps(tmps)
|
||||
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Errorf("cleanupTemps did not remove %q (err=%v)", path, err)
|
||||
}
|
||||
}
|
||||
|
||||
// A non-positive limit must fall back to the 15000 default rather than
|
||||
// spilling every non-empty result.
|
||||
func TestOffloadToolResultZeroLimitUsesDefault(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
var tmps []string
|
||||
res := strings.Repeat("y", 15000) // exactly the default limit: stays inline
|
||||
if got := offloadToolResult(res, 0, &tmps); got != res {
|
||||
t.Errorf("result at default limit should stay inline")
|
||||
}
|
||||
if len(tmps) != 0 {
|
||||
t.Errorf("no spill expected at the limit, got %v", tmps)
|
||||
}
|
||||
}
|
||||
|
||||
// End-to-end: AL() offloads an oversized shell result, hands the model the
|
||||
// instruction instead of the raw bytes, and deletes the spilled file once the
|
||||
// round finishes.
|
||||
func TestALOffloadsLargeToolResultAndCleans(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
t.Setenv("TMPDIR", tmp)
|
||||
seen := ""
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var req struct {
|
||||
Messages []Message `json:"messages"`
|
||||
}
|
||||
json.NewDecoder(r.Body).Decode(&req)
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "tool" && m.Content != nil {
|
||||
seen = *m.Content
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"done"}}]}`))
|
||||
return
|
||||
}
|
||||
}
|
||||
args, _ := json.Marshal(map[string]any{"command": "yes x | head -c 20000"})
|
||||
w.Write([]byte(fmt.Sprintf(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"shell_exec","arguments":%s}}]}}]}`, strconv.Quote(string(args)))))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := defCfg
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
cfg.MaxToolRes = 15000
|
||||
_, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
if seen == "" {
|
||||
t.Fatalf("no tool message reached the model")
|
||||
}
|
||||
if len(seen) > cfg.MaxToolRes {
|
||||
t.Errorf("tool message sent to the model is %d bytes, exceeds limit %d", len(seen), cfg.MaxToolRes)
|
||||
}
|
||||
if !strings.Contains(seen, "too large for context") {
|
||||
t.Errorf("expected an overflow instruction, got %q", seen)
|
||||
}
|
||||
// The spilled file must have been deleted when the round ended.
|
||||
dir := filepath.Join(tmp, "bantam", "toolres")
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
t.Fatalf("read toolres dir: %v", err)
|
||||
}
|
||||
if len(entries) != 0 {
|
||||
t.Errorf("expected no leftover temp files, found %d", len(entries))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// buildSystemPrompt must substitute the configured max_tool_res value and the
|
||||
// concrete temporary directory into the built-in prompt.
|
||||
func TestBuildSystemPromptMaxToolResHint(t *testing.T) {
|
||||
t.Setenv("TMPDIR", "/custom/tmp")
|
||||
cfg := defCfg
|
||||
cfg.MaxToolRes = 12345
|
||||
sp := buildSystemPrompt(&cfg)
|
||||
if !strings.Contains(sp, "12345") {
|
||||
t.Errorf("prompt should embed the configured max_tool_res (12345): %q", sp)
|
||||
}
|
||||
if strings.Contains(sp, "$MAX_TOOL_RES") {
|
||||
t.Errorf("prompt still contains the unexpanded $MAX_TOOL_RES placeholder: %q", sp)
|
||||
}
|
||||
if !strings.Contains(sp, filepath.Join("/custom/tmp", "bantam", "toolres")) {
|
||||
t.Errorf("prompt should reference the concrete toolres dir: %q", sp)
|
||||
}
|
||||
if strings.Contains(sp, "$TMPDIR/bantam") {
|
||||
t.Errorf("prompt still contains the unexpanded $TMPDIR/bantam placeholder: %q", sp)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user