148 lines · 5.2 KB · Swift
RawDownload
import Foundation extension CostUsageScanner {    static func extractCodexTurnContextModel(from bytes: Data) -> String? {        guard let text = truncatedUTF8String(from: bytes) else { return nil }        let object = text[...]        guard Self.extractJSONStringField("type", from: object, atDepth: 1) == "turn_context",              let payloadText = Self.extractJSONObjectField("payload", from: object, atDepth: 1)        else { return nil }         let payloadModel = Self.extractJSONStringField("model", from: payloadText, atDepth: 1)            ?? Self.extractJSONStringField("model_name", from: payloadText, atDepth: 1)        if let payloadModel { return payloadModel }         guard let infoText = Self.extractJSONObjectField("info", from: payloadText, atDepth: 1) else { return nil }        return Self.extractJSONStringField("model", from: infoText, atDepth: 1)            ?? Self.extractJSONStringField("model_name", from: infoText, atDepth: 1)    }     static func truncatedUTF8String(from bytes: Data) -> String? {        for dropCount in 0...min(4, bytes.count) {            let end = bytes.count - dropCount            if let text = String(bytes: bytes.prefix(end), encoding: .utf8) {                return text            }        }        return nil    }     static func extractJSONStringField(        _ field: String,        from text: Substring,        atDepth targetDepth: Int) -> String?    {        self.extractJSONField(field, from: text, atDepth: targetDepth) { text, index in            guard index < text.endIndex, text[index] == "\"" else { return nil }            let value = Self.parseJSONString(in: text, index: &index)            return value?.isEmpty == true ? nil : value        }    }     static func extractJSONObjectField(        _ field: String,        from text: Substring,        atDepth targetDepth: Int) -> Substring?    {        self.extractJSONField(field, from: text, atDepth: targetDepth) { text, index in            guard index < text.endIndex, text[index] == "{" else { return nil }            return text[index...]        }    }     static func extractJSONIntField(        _ field: String,        from text: Substring,        atDepth targetDepth: Int) -> Int?    {        self.extractJSONField(field, from: text, atDepth: targetDepth) { text, index in            Self.parseJSONInt(in: text, index: &index)        }    }     static func extractJSONField<T>(        _ field: String,        from text: Substring,        atDepth targetDepth: Int,        parseValue: (Substring, inout String.Index) -> T?) -> T?    {        var index = text.startIndex        var depth = 0         while index < text.endIndex {            let character = text[index]            if character == "{" {                depth += 1                text.formIndex(after: &index)            } else if character == "}" {                depth -= 1                text.formIndex(after: &index)            } else if character == "\"" {                var valueIndex = index                guard let key = Self.parseJSONString(in: text, index: &valueIndex) else { return nil }                defer { index = valueIndex }                guard depth == targetDepth, key == field else { continue }                 Self.skipJSONWhitespace(in: text, index: &valueIndex)                guard valueIndex < text.endIndex, text[valueIndex] == ":" else { continue }                 text.formIndex(after: &valueIndex)                Self.skipJSONWhitespace(in: text, index: &valueIndex)                if let value = parseValue(text, &valueIndex) {                    return value                }            } else {                text.formIndex(after: &index)            }        }         return nil    }     static func parseJSONString(in text: Substring, index: inout String.Index) -> String? {        guard index < text.endIndex, text[index] == "\"" else { return nil }        text.formIndex(after: &index)        var value = ""        var isEscaped = false        while index < text.endIndex {            let character = text[index]            if isEscaped {                value.append(character)                isEscaped = false            } else if character == "\\" {                isEscaped = true            } else if character == "\"" {                text.formIndex(after: &index)                return value            } else {                value.append(character)            }            text.formIndex(after: &index)        }         return nil    }     static func parseJSONInt(in text: Substring, index: inout String.Index) -> Int? {        var sign = 1        if index < text.endIndex, text[index] == "-" {            sign = -1            text.formIndex(after: &index)        }         var value = 0        var sawDigit = false        while index < text.endIndex, let digit = text[index].wholeNumberValue {            sawDigit = true            value = (value * 10) + digit            text.formIndex(after: &index)        }        return sawDigit ? value * sign : nil    }     static func skipJSONWhitespace(in text: Substring, index: inout String.Index) {        while index < text.endIndex, text[index].isWhitespace {            text.formIndex(after: &index)        }    }}