openteam/server/internal/proxy/jsonfast.go

package proxy

import (
	"bytes"
	"encoding/json"
)

// 请求体上的定点改写：只动需要改的那一段，不做整包 unmarshal → map → marshal。
// 整包往返在 agent 场景（上下文几十到几百 KB）实测约 0.5–0.6ms 且分配翻倍，
// 而这里只做一次 O(n) 扫描。

// rewriteModelField 改写请求体顶层 model 字段，其余字节原样保留。
// 嵌套出现的 "model"（比如消息内容里引用的 JSON）不会被误改：只认顶层键。
// 找不到顶层 model 或请求体不是对象时返回 (body, false)。
func rewriteModelField(body []byte, model string) ([]byte, bool) {
	span, ok := topLevelValueSpan(body, "model")
	if !ok {
		return body, false
	}
	quoted, err := json.Marshal(model)
	if err != nil {
		return body, false
	}
	out := make([]byte, 0, len(body)-(span[1]-span[0])+len(quoted))
	out = append(out, body[:span[0]]...)
	out = append(out, quoted...)
	out = append(out, body[span[1]:]...)
	return out, true
}

// injectStreamOptions 在请求体顶层追加 stream_options，让流式末块带上 usage。
// 只在没有该键时插入；body 必须是 JSON 对象（以 '}' 结尾，允许尾部空白）。
func injectStreamOptions(body []byte, options string) ([]byte, bool) {
	trimmed := bytes.TrimRight(body, " \t\r\n")
	if len(trimmed) < 2 || trimmed[0] != '{' || trimmed[len(trimmed)-1] != '}' {
		return body, false
	}
	if bytes.Contains(trimmed, []byte(`"stream_options"`)) {
		return body, false
	}
	end := len(trimmed) - 1
	out := make([]byte, 0, len(body)+len(options)+16)
	out = append(out, trimmed[:end]...)
	// 空对象 {} 不能加逗号
	if last := bytes.TrimRight(trimmed[:end], " \t\r\n"); len(last) > 1 {
		out = append(out, ',')
	}
	out = append(out, `"stream_options":`...)
	out = append(out, options...)
	out = append(out, '}')
	return out, true
}

// topLevelValueSpan 找顶层键 key 的值在 body 中的字节区间 [start, end)。
// 只遍历顶层（depth==1）的键，跳过字符串字面量与转义。
func topLevelValueSpan(body []byte, key string) ([2]int, bool) {
	var zero [2]int
	trimmed := bytes.TrimLeft(body, " \t\r\n")
	if len(trimmed) == 0 || trimmed[0] != '{' {
		return zero, false
	}
	// 重新对齐到 trimmed 起点后的偏移
	base := len(body) - len(trimmed)
	keyPat := []byte(`"` + key + `"`)
	i := base + 1
	depth := 1
	for i < len(body) {
		switch body[i] {
		case '{', '[':
			depth++
			i++
		case '}', ']':
			depth--
			if depth == 0 {
				return zero, false
			}
			i++
		case '"':
			start := i
			i = skipJSONString(body, i)
			// 顶层字符串：可能是"键"，也可能是"值"
			if depth == 1 && bytes.Equal(body[start:i], keyPat) {
				j := i
				for j < len(body) && (body[j] == ' ' || body[j] == '\t' || body[j] == '\r' || body[j] == '\n') {
					j++
				}
				if j < len(body) && body[j] == ':' {
					j++
					for j < len(body) && (body[j] == ' ' || body[j] == '\t' || body[j] == '\r' || body[j] == '\n') {
						j++
					}
					vs := j
					ve, ok := jsonValueEnd(body, vs)
					if !ok {
						return zero, false
					}
					return [2]int{vs, ve}, true
				}
			}
		default:
			i++
		}
	}
	return zero, false
}

// skipJSONString 返回从 body[i]（'"'）开始的下一个未转义引号之后的位置。
func skipJSONString(body []byte, i int) int {
	i++ // 跳过开引号
	for i < len(body) {
		switch body[i] {
		case '\\':
			i += 2
		case '"':
			return i + 1
		default:
			i++
		}
	}
	return i
}

// jsonValueEnd 返回从 i 开始的一个 JSON 值的结束位置（不含尾部空白）。
func jsonValueEnd(body []byte, i int) (int, bool) {
	if i >= len(body) {
		return 0, false
	}
	switch body[i] {
	case '"':
		return skipJSONString(body, i), true
	case '{', '[':
		depth := 0
		for j := i; j < len(body); j++ {
			switch body[j] {
			case '"':
				j = skipJSONString(body, j) - 1
			case '{', '[':
				depth++
			case '}', ']':
				depth--
				if depth == 0 {
					return j + 1, true
				}
			}
		}
		return 0, false
	default:
		j := i
		for j < len(body) && body[j] != ',' && body[j] != '}' && body[j] != ']' &&
			body[j] != ' ' && body[j] != '\t' && body[j] != '\r' && body[j] != '\n' {
			j++
		}
		return j, j > i
	}
}