IPI protection

main.go

@@ -22,6 +22,7 @@ import (
   "strconv"
   "strings"
   "time"
+  "unicode"
   "unicode/utf8"
 )
 
@@ -712,11 +713,27 @@ func cleanMessagesForLLM(msgs []Message) []Message {
   return out
 }
 
+func filterText(s string) string {
+  var b strings.Builder
+  b.Grow(len(s))
+  for _, r := range s {
+    if r == ' ' || r == '\t' || r == '\n' {
+      b.WriteRune(r)
+    } else if unicode.Is(unicode.Z, r) || unicode.IsControl(r) || unicode.Is(unicode.C, r) {
+      continue
+    } else if unicode.IsPrint(r) {
+      b.WriteRune(r)
+    }
+  }
+  return b.String()
+}
+
 func sanitizeMessages(msgs []Message) {
   for i := range msgs {
     if msgs[i].Role == "assistant" && len(msgs[i].ToolCalls) > 0 {
       for j := range msgs[i].ToolCalls {
         tc := &msgs[i].ToolCalls[j]
+        tc.Function.Arguments = filterText(tc.Function.Arguments)
         astr := tc.Function.Arguments
         var a map[string]any
         if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil {
@@ -724,6 +741,8 @@ func sanitizeMessages(msgs []Message) {
           tc.Function.Arguments = string(fixed)
         }
       }
+    } else if msgs[i].Role == "tool" && msgs[i].Content != nil {
+      msgs[i].Content = strp(filterText(*msgs[i].Content))
     }
   }
 }
@@ -956,12 +975,13 @@ func parseStream(ctx context.Context, r io.Reader) (Message, Usage, error) {
 }
 
 func shell(ctx context.Context, cmd string, timeout int) string {
+  cmd = filterText(cmd)
   cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
   defer cancel()
   c := exec.CommandContext(cmdCtx, "sh", "-c", cmd)
   c.WaitDelay = 100 * time.Millisecond
   out, err := c.CombinedOutput()
-  res := strings.TrimSpace(string(out))
+  res := strings.TrimSpace(filterText(string(out)))
   if ctx.Err() != nil {
     return "[interrupted]\n\nexit: -1"
   }
@@ -1018,6 +1038,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
     if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() }
     for j := range m.ToolCalls {
       tc := &m.ToolCalls[j]
+      tc.Function.Arguments = filterText(tc.Function.Arguments)
       astr := tc.Function.Arguments
       var a map[string]any
       if err := json.Unmarshal([]byte(astr), &a); err != nil || a == nil {
@@ -1035,7 +1056,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
     if len(m.ToolCalls) == 0 { done = true; break }
     for _, tc := range m.ToolCalls {
       if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
-      fn, astr := tc.Function.Name, tc.Function.Arguments
+      fn, astr := tc.Function.Name, filterText(tc.Function.Arguments)
       fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33))
       res, sty := "", 2
       var a map[string]any
@@ -1045,10 +1066,12 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
         switch fn {
         case "shell_exec":
           cmd, _ := a["command"].(string)
+          cmd = filterText(cmd)
           res = shell(ctx, cmd, cfg.ShellTimeout)
           if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
         case "run_subagent":
           pr, _ := a["prompt"].(string)
+          pr = filterText(pr)
           if depth >= MAX_DEPTH {
             res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31
           } else {
@@ -1064,13 +1087,14 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
             } else {
               turnUsage.CompletionTokens += subu.CompletionTokens
               turnUsage.TotalTokens += subu.TotalTokens
-              res, sty = last(subr), 2
+              res, sty = filterText(last(subr)), 2
             }
           }
         default:
           res, sty = "Unknown tool: " + fn, 31
         }
       }
+      res = filterText(res)
       fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n")
       msgs = append(msgs, Message{Role: "tool", ToolCallID: tc.ID, Content: strp(res)})
     }
@@ -1422,7 +1446,7 @@ func main() {
     }
     u := strings.TrimSpace(string(data))
     if strings.HasPrefix(u, "!") {
-      cmd := strings.TrimSpace(strings.TrimPrefix(u, "!"))
+      cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!")))
       if cmd != "" {
         astr, _ := json.Marshal(map[string]string{"command": cmd})
         fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
@@ -1551,7 +1575,7 @@ func main() {
       }
       continue
     case strings.HasPrefix(u, "!"):
-      cmd := strings.TrimSpace(strings.TrimPrefix(u, "!"))
+      cmd := filterText(strings.TrimSpace(strings.TrimPrefix(u, "!")))
       if cmd != "" {
         astr, _ := json.Marshal(map[string]string{"command": cmd})
         fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))