magpie-gemini-web 0.3.3 → 0.4.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.en.md +9 -3
- package/README.md +9 -3
- package/index.mjs +15 -1
- package/lib/backend.mjs +5 -3
- package/lib/task.mjs +87 -0
- package/native-task/LICENSE +21 -0
- package/native-task/README.md +87 -0
- package/native-task/bridge/gemini_native_bridge/__init__.py +3 -0
- package/native-task/bridge/gemini_native_bridge/__main__.py +235 -0
- package/native-task/bridge/gemini_native_bridge/auth.py +143 -0
- package/native-task/bridge/gemini_native_bridge/config.py +222 -0
- package/native-task/bridge/gemini_native_bridge/contract.py +283 -0
- package/native-task/bridge/gemini_native_bridge/native_auth_keychain.swift +111 -0
- package/native-task/bridge/gemini_native_bridge/oauth.py +154 -0
- package/native-task/bridge/gemini_native_bridge/protocol.py +363 -0
- package/native-task/bridge/gemini_native_bridge/server.py +441 -0
- package/native-task/bridge/gemini_native_bridge/state.py +806 -0
- package/native-task/bridge/gemini_native_bridge/transport.py +531 -0
- package/native-task/index.mjs +76 -0
- package/native-task/lib/runtime.mjs +189 -0
- package/native-task/pyproject.toml +18 -0
- package/native-task/requirements.txt +2 -0
- package/native-task/scripts/bridge.mjs +97 -0
- package/package.json +11 -3
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
// Reads only Gemini Safe Storage through the normal, noninteractive Keychain API,
|
|
2
|
+
// then decrypts one caller-selected Gemini auth file. No Frida or App process.
|
|
3
|
+
import Foundation
|
|
4
|
+
import Security
|
|
5
|
+
import LocalAuthentication
|
|
6
|
+
import CryptoKit
|
|
7
|
+
import Darwin
|
|
8
|
+
|
|
9
|
+
enum ProbeError: Error { case invalidFormat, keychain(Int32), decryption, output }
|
|
10
|
+
struct Field { let number: Int; let integer: UInt64?; let data: Data? }
|
|
11
|
+
func fields(_ input: Data) throws -> [Field] {
|
|
12
|
+
let bytes = Array(input); var pos = 0; var result = [Field]()
|
|
13
|
+
func varint() throws -> UInt64 {
|
|
14
|
+
var value: UInt64 = 0
|
|
15
|
+
for shift in stride(from: 0, to: 70, by: 7) {
|
|
16
|
+
guard pos < bytes.count else { throw ProbeError.invalidFormat }
|
|
17
|
+
let byte = bytes[pos]; pos += 1
|
|
18
|
+
if shift == 63 && byte > 1 { throw ProbeError.invalidFormat }
|
|
19
|
+
value |= UInt64(byte & 127) << shift
|
|
20
|
+
if byte < 128 { return value }
|
|
21
|
+
}
|
|
22
|
+
throw ProbeError.invalidFormat
|
|
23
|
+
}
|
|
24
|
+
while pos < bytes.count {
|
|
25
|
+
let key = try varint(); let number = Int(key >> 3)
|
|
26
|
+
guard number > 0 else { throw ProbeError.invalidFormat }
|
|
27
|
+
if key & 7 == 0 {
|
|
28
|
+
result.append(Field(number: number, integer: try varint(), data: nil))
|
|
29
|
+
} else if key & 7 == 2 {
|
|
30
|
+
let size = try varint()
|
|
31
|
+
guard size <= UInt64(bytes.count - pos) else { throw ProbeError.invalidFormat }
|
|
32
|
+
result.append(Field(number: number, integer: nil, data: Data(bytes[pos..<pos+Int(size)])))
|
|
33
|
+
pos += Int(size)
|
|
34
|
+
} else { throw ProbeError.invalidFormat }
|
|
35
|
+
}
|
|
36
|
+
return result
|
|
37
|
+
}
|
|
38
|
+
func one(_ fs: [Field], _ number: Int) throws -> Field {
|
|
39
|
+
let found = fs.filter { $0.number == number }
|
|
40
|
+
guard found.count == 1 else { throw ProbeError.invalidFormat }
|
|
41
|
+
return found[0]
|
|
42
|
+
}
|
|
43
|
+
func save(_ path: String, _ data: Data) throws {
|
|
44
|
+
let fd = open(path, O_WRONLY | O_CREAT | O_EXCL, mode_t(0o600))
|
|
45
|
+
guard fd >= 0 else { throw ProbeError.output }
|
|
46
|
+
defer { close(fd) }
|
|
47
|
+
var offset = 0
|
|
48
|
+
try data.withUnsafeBytes { bytes in
|
|
49
|
+
while offset < bytes.count {
|
|
50
|
+
let n = write(fd, bytes.baseAddress!.advanced(by: offset), bytes.count-offset)
|
|
51
|
+
guard n > 0 else { throw ProbeError.output }
|
|
52
|
+
offset += n
|
|
53
|
+
}
|
|
54
|
+
}
|
|
55
|
+
}
|
|
56
|
+
var report: [String: Any] = ["success": false, "interaction_allowed": false]
|
|
57
|
+
do {
|
|
58
|
+
guard CommandLine.arguments.count == 3 else { throw ProbeError.invalidFormat }
|
|
59
|
+
let context = LAContext(); context.interactionNotAllowed = true
|
|
60
|
+
let query: [String: Any] = [
|
|
61
|
+
kSecClass as String: kSecClassGenericPassword,
|
|
62
|
+
kSecAttrService as String: "Gemini Safe Storage",
|
|
63
|
+
kSecMatchLimit as String: kSecMatchLimitOne,
|
|
64
|
+
kSecReturnData as String: true,
|
|
65
|
+
kSecUseAuthenticationContext as String: context,
|
|
66
|
+
]
|
|
67
|
+
var item: CFTypeRef?
|
|
68
|
+
let status = SecItemCopyMatching(query as CFDictionary, &item)
|
|
69
|
+
report["keychain_status"] = status
|
|
70
|
+
guard status == errSecSuccess, let stored = item as? Data else { throw ProbeError.keychain(status) }
|
|
71
|
+
let fs = try fields(stored)
|
|
72
|
+
guard try one(fs, 1).integer == 3735928558 else { throw ProbeError.invalidFormat }
|
|
73
|
+
let keys = try fs.filter { $0.number == 2 }.map { field -> Data in
|
|
74
|
+
guard let data = field.data else { throw ProbeError.invalidFormat }
|
|
75
|
+
let entry = try fields(data)
|
|
76
|
+
guard let key = try one(entry, 4).data, key.count == 32 else { throw ProbeError.invalidFormat }
|
|
77
|
+
return key
|
|
78
|
+
}
|
|
79
|
+
guard !keys.isEmpty && keys.count <= 16 else { throw ProbeError.invalidFormat }
|
|
80
|
+
let encrypted = try Data(contentsOf: URL(fileURLWithPath: CommandLine.arguments[1]))
|
|
81
|
+
let bytes = Array(encrypted)
|
|
82
|
+
guard bytes.count > 28 && bytes.count < 1024*1024,
|
|
83
|
+
Array(bytes[0..<8]) == [0x2e,0xba,0xfe,0xca,0x2e,0xba,0xfe,0xca] else { throw ProbeError.invalidFormat }
|
|
84
|
+
let length = (0..<4).reduce(0) { $0 | Int(bytes[8+$1]) << (8*$1) }
|
|
85
|
+
guard length > 0 && length < 65536 && 12+length < bytes.count else { throw ProbeError.invalidFormat }
|
|
86
|
+
let header = try fields(Data(bytes[12..<12+length]))
|
|
87
|
+
let body = Data(bytes[(12+length)...])
|
|
88
|
+
guard try one(header, 1).integer == 1, try one(header, 2).integer == 1,
|
|
89
|
+
try one(header, 5).integer == UInt64(body.count), body.count > 16,
|
|
90
|
+
let nonce = try one(header, 4).data, nonce.count == 12 else { throw ProbeError.invalidFormat }
|
|
91
|
+
let box = try AES.GCM.SealedBox(nonce: AES.GCM.Nonce(data: nonce),
|
|
92
|
+
ciphertext: body.dropLast(16), tag: body.suffix(16))
|
|
93
|
+
var plaintext: Data? = nil
|
|
94
|
+
for key in keys {
|
|
95
|
+
if let decoded = try? AES.GCM.open(box, using: SymmetricKey(data: key)) {
|
|
96
|
+
guard plaintext == nil else { throw ProbeError.invalidFormat }
|
|
97
|
+
plaintext = decoded
|
|
98
|
+
}
|
|
99
|
+
}
|
|
100
|
+
guard let plain = plaintext else { throw ProbeError.decryption }
|
|
101
|
+
_ = try fields(plain)
|
|
102
|
+
try save(CommandLine.arguments[2], plain)
|
|
103
|
+
report["success"] = true
|
|
104
|
+
report["plaintext_bytes"] = plain.count
|
|
105
|
+
report["decryption"] = "AES-256-GCM"
|
|
106
|
+
} catch {
|
|
107
|
+
report["error_type"] = String(describing: error)
|
|
108
|
+
}
|
|
109
|
+
let result = try! JSONSerialization.data(withJSONObject: report, options: [.sortedKeys])
|
|
110
|
+
print(String(data: result, encoding: .utf8)!)
|
|
111
|
+
exit(report["success"] as? Bool == true ? 0 : 2)
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
"""Existing-login credential primitives; fixed Google endpoints, no secret logging."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
import base64
|
|
5
|
+
import hashlib
|
|
6
|
+
import http.client
|
|
7
|
+
import json
|
|
8
|
+
import struct
|
|
9
|
+
import subprocess
|
|
10
|
+
import tempfile
|
|
11
|
+
import time
|
|
12
|
+
import urllib.error
|
|
13
|
+
import urllib.parse
|
|
14
|
+
import urllib.request
|
|
15
|
+
from dataclasses import dataclass
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
|
|
18
|
+
from .contract import BridgeError
|
|
19
|
+
from .protocol import BINARY_SHA256, fields, one
|
|
20
|
+
|
|
21
|
+
TOKEN_URL = "https://www.googleapis.com/oauth2/v4/token"
|
|
22
|
+
TOKENINFO_URL = "https://oauth2.googleapis.com/tokeninfo"
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@dataclass(frozen=True)
|
|
26
|
+
class Credential:
|
|
27
|
+
subject: str
|
|
28
|
+
email: str
|
|
29
|
+
refresh_token: str
|
|
30
|
+
audience: str
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class NoRedirect(urllib.request.HTTPRedirectHandler):
|
|
34
|
+
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
35
|
+
return None
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _json_request(url: str, data: bytes | None = None):
|
|
39
|
+
if url != TOKEN_URL and not url.startswith(TOKENINFO_URL + "?"):
|
|
40
|
+
raise BridgeError("unexpected_auth_endpoint", "Authentication endpoint is not allowed.")
|
|
41
|
+
request = urllib.request.Request(url, data=data, headers={
|
|
42
|
+
"Content-Type": "application/x-www-form-urlencoded"} if data is not None else {})
|
|
43
|
+
opener = urllib.request.build_opener(NoRedirect)
|
|
44
|
+
for attempt in range(2):
|
|
45
|
+
try:
|
|
46
|
+
with opener.open(request, timeout=15) as response:
|
|
47
|
+
raw = response.read(65537)
|
|
48
|
+
if len(raw) > 65536:
|
|
49
|
+
raise BridgeError("invalid_auth_response", "Authentication response exceeds the allowed size.")
|
|
50
|
+
body = json.loads(raw)
|
|
51
|
+
if not isinstance(body, dict):
|
|
52
|
+
raise ValueError()
|
|
53
|
+
return response.status, body
|
|
54
|
+
except urllib.error.HTTPError as error:
|
|
55
|
+
try:
|
|
56
|
+
body = json.loads(error.read(65536))
|
|
57
|
+
reason = body.get("error") if isinstance(body, dict) else None
|
|
58
|
+
except (ValueError, OSError, http.client.HTTPException):
|
|
59
|
+
reason = None
|
|
60
|
+
return error.code, {"error": reason if reason in {"invalid_grant", "invalid_token", "invalid_client"} else "rejected"}
|
|
61
|
+
except (http.client.HTTPException, urllib.error.URLError, OSError):
|
|
62
|
+
if attempt:
|
|
63
|
+
raise BridgeError("auth_transport_failed", "Cannot reach the Google authentication service.") from None
|
|
64
|
+
except (ValueError, UnicodeError):
|
|
65
|
+
raise BridgeError("invalid_auth_response", "Google returned an invalid authentication response.") from None
|
|
66
|
+
raise AssertionError("unreachable")
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def token_info(token: str):
|
|
70
|
+
return _json_request(TOKENINFO_URL + "?" + urllib.parse.urlencode({"access_token": token}))
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def extract_client(binary: Path) -> tuple[str, str]:
|
|
74
|
+
data = binary.read_bytes()
|
|
75
|
+
if hashlib.sha256(data).hexdigest() != BINARY_SHA256:
|
|
76
|
+
raise BridgeError("unsupported_app_version", "This bridge requires the validated Gemini macOS 1.128.14.1166 binary.")
|
|
77
|
+
if struct.unpack_from("<I", data)[0] != 0xfeedfacf:
|
|
78
|
+
raise BridgeError("unsupported_app_binary", "Expected the validated arm64 Mach-O binary.")
|
|
79
|
+
segments = []
|
|
80
|
+
p = 32
|
|
81
|
+
for _ in range(struct.unpack_from("<I", data, 16)[0]):
|
|
82
|
+
cmd, size = struct.unpack_from("<II", data, p)
|
|
83
|
+
if cmd == 0x19:
|
|
84
|
+
vmaddr, vmsize, offset, length = struct.unpack_from("<QQQQ", data, p + 24)
|
|
85
|
+
segments.append((vmaddr, offset, length))
|
|
86
|
+
p += size
|
|
87
|
+
|
|
88
|
+
def address(va, length):
|
|
89
|
+
offsets = [off + va - vm for vm, off, size in segments if vm <= va and va + length <= vm + size]
|
|
90
|
+
if len(offsets) != 1:
|
|
91
|
+
raise BridgeError("invalid_client_config", "OAuth client configuration is outside the validated segment.")
|
|
92
|
+
return data[offsets[0]:offsets[0] + length]
|
|
93
|
+
|
|
94
|
+
client_id = address(0x109e17220, 72).decode()
|
|
95
|
+
client_secret = bytes(v ^ 0x5a for v in address(0x10aa4b380, 35)).decode()
|
|
96
|
+
if not client_id.endswith(".apps.googleusercontent.com") or not client_secret.startswith("GOCSPX-"):
|
|
97
|
+
raise BridgeError("invalid_client_config", "The installed OAuth client configuration is not recognized.")
|
|
98
|
+
return client_id, client_secret
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def read_credential(auth_file: Path, work_dir: Path) -> Credential:
|
|
102
|
+
helper = Path(__file__).with_name("native_auth_keychain.swift")
|
|
103
|
+
try:
|
|
104
|
+
with tempfile.TemporaryDirectory(prefix="credential-", dir=work_dir) as directory:
|
|
105
|
+
output = Path(directory) / "plain.bin"
|
|
106
|
+
result = subprocess.run(["swift", str(helper), str(auth_file), str(output)],
|
|
107
|
+
capture_output=True, timeout=45)
|
|
108
|
+
if result.returncode:
|
|
109
|
+
raise BridgeError("keychain_unavailable", "Cannot read this Gemini login through the noninteractive Keychain API.")
|
|
110
|
+
data = output.read_bytes()
|
|
111
|
+
subject = one(data, 1).decode()
|
|
112
|
+
email = one(data, 2).decode()
|
|
113
|
+
refresh = one(data, 4).decode()
|
|
114
|
+
encoded = one(data, 8).split(b".")[1]
|
|
115
|
+
claims = json.loads(base64.urlsafe_b64decode(encoded + b"=" * (-len(encoded) % 4)))
|
|
116
|
+
if not subject or claims.get("sub") != subject or not refresh.startswith("1//"):
|
|
117
|
+
raise ValueError()
|
|
118
|
+
audience = claims.get("aud")
|
|
119
|
+
if not isinstance(audience, str) or not audience.endswith(".apps.googleusercontent.com"):
|
|
120
|
+
raise ValueError()
|
|
121
|
+
return Credential(subject, email, refresh, audience)
|
|
122
|
+
except BridgeError:
|
|
123
|
+
raise
|
|
124
|
+
except (OSError, ValueError, IndexError, subprocess.SubprocessError):
|
|
125
|
+
raise BridgeError("credential_unavailable", "The selected local Gemini login could not be decoded.") from None
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def refresh_access(client: tuple[str, str], credential: Credential, rotated_refresh: str | None = None) -> dict:
|
|
129
|
+
client_id, client_secret = client
|
|
130
|
+
if client_id != credential.audience:
|
|
131
|
+
raise BridgeError("oauth_client_mismatch", "The local login belongs to a different OAuth client.")
|
|
132
|
+
refresh = rotated_refresh or credential.refresh_token
|
|
133
|
+
form = urllib.parse.urlencode({"grant_type": "refresh_token", "client_id": client_id,
|
|
134
|
+
"client_secret": client_secret, "refresh_token": refresh}).encode()
|
|
135
|
+
status, body = _json_request(TOKEN_URL, form)
|
|
136
|
+
if status != 200 or not isinstance(body.get("access_token"), str) or body.get("token_type", "").lower() != "bearer":
|
|
137
|
+
raise BridgeError("login_refresh_failed", "The Gemini login cannot be refreshed; sign in through the app again if it has expired.")
|
|
138
|
+
token = body["access_token"]
|
|
139
|
+
status, info = token_info(token)
|
|
140
|
+
subject = str(info.get("sub", info.get("user_id", "")))
|
|
141
|
+
if status != 200 or subject != credential.subject:
|
|
142
|
+
raise BridgeError("account_mismatch", "The renewed token does not match the configured account.")
|
|
143
|
+
try:
|
|
144
|
+
expires = min(float(info["exp"]), time.time() + int(body["expires_in"]))
|
|
145
|
+
if not time.time() + 60 < expires <= time.time() + 86400:
|
|
146
|
+
raise ValueError()
|
|
147
|
+
except (KeyError, ValueError, TypeError, OverflowError):
|
|
148
|
+
raise BridgeError("invalid_token_lifetime", "The renewed token has an invalid expiry.") from None
|
|
149
|
+
rotation = body.get("refresh_token")
|
|
150
|
+
if rotation is not None and (not isinstance(rotation, str) or not rotation.startswith("1//")):
|
|
151
|
+
raise BridgeError("invalid_refresh_rotation", "Google returned an unsupported refresh-token rotation.")
|
|
152
|
+
return {"access_token": token, "expires_at": expires,
|
|
153
|
+
"subject_sha256": hashlib.sha256(subject.encode()).hexdigest(),
|
|
154
|
+
"rotated_refresh_token": rotation or rotated_refresh}
|
|
@@ -0,0 +1,363 @@
|
|
|
1
|
+
"""Validated Gemini 1.128.14.1166 wire paths. No networking or credential I/O."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
import math
|
|
5
|
+
import struct
|
|
6
|
+
import uuid
|
|
7
|
+
from dataclasses import dataclass
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
HOST = "robinfrontend-pa.googleapis.com"
|
|
11
|
+
TASK_RPC = "/assistant.mobile.robin.RobinService/ProcessQuery"
|
|
12
|
+
TOOLS_RPC = "/assistant.mobile.robin.notify.GeminiNotifyService/ExecuteRemoteTools"
|
|
13
|
+
TURN_RPC = "/assistant.mobile.robin.RobinService/ListConversationTurns"
|
|
14
|
+
MODEL_ID = "gemini-task-argon"
|
|
15
|
+
MODE_ID = "95a4c0abcba24364"
|
|
16
|
+
SERVER_NAME = "agent_bridge"
|
|
17
|
+
BINARY_SHA256 = "8cd89234b23182234e8e211e1449ae44e0a831b2185e98f4d1371a83c66f8964"
|
|
18
|
+
MAX_FRAME = 8 * 1024 * 1024
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class ProtocolError(ValueError):
|
|
22
|
+
pass
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def varint(data: bytes, offset: int) -> tuple[int, int]:
|
|
26
|
+
value = 0
|
|
27
|
+
for shift in range(0, 70, 7):
|
|
28
|
+
if offset >= len(data):
|
|
29
|
+
raise ProtocolError("truncated_varint")
|
|
30
|
+
byte = data[offset]
|
|
31
|
+
offset += 1
|
|
32
|
+
if shift == 63 and byte > 1:
|
|
33
|
+
raise ProtocolError("varint_overflow")
|
|
34
|
+
value |= (byte & 127) << shift
|
|
35
|
+
if byte < 128:
|
|
36
|
+
return value, offset
|
|
37
|
+
raise ProtocolError("invalid_varint")
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def encode_varint(value: int) -> bytes:
|
|
41
|
+
if not isinstance(value, int) or value < 0 or value >= 1 << 64:
|
|
42
|
+
raise ProtocolError("invalid_unsigned_integer")
|
|
43
|
+
out = bytearray()
|
|
44
|
+
while value > 127:
|
|
45
|
+
out.append((value & 127) | 128)
|
|
46
|
+
value >>= 7
|
|
47
|
+
out.append(value)
|
|
48
|
+
return bytes(out)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def fields(data: bytes):
|
|
52
|
+
result = []
|
|
53
|
+
pos = 0
|
|
54
|
+
while pos < len(data):
|
|
55
|
+
start = pos
|
|
56
|
+
key, pos = varint(data, pos)
|
|
57
|
+
number, wire = key >> 3, key & 7
|
|
58
|
+
if not 0 < number < 1 << 29:
|
|
59
|
+
raise ProtocolError("invalid_field_number")
|
|
60
|
+
if wire == 0:
|
|
61
|
+
value, pos = varint(data, pos)
|
|
62
|
+
elif wire in (1, 5, 2):
|
|
63
|
+
if wire == 2:
|
|
64
|
+
size, pos = varint(data, pos)
|
|
65
|
+
else:
|
|
66
|
+
size = 8 if wire == 1 else 4
|
|
67
|
+
if size > len(data) - pos:
|
|
68
|
+
raise ProtocolError("truncated_field")
|
|
69
|
+
value = data[pos:pos + size]
|
|
70
|
+
pos += size
|
|
71
|
+
else:
|
|
72
|
+
raise ProtocolError("unsupported_wire_type")
|
|
73
|
+
result.append((number, wire, value, data[start:pos]))
|
|
74
|
+
return result
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def one(data: bytes, number: int):
|
|
78
|
+
found = [v for n, _, v, _ in fields(data) if n == number]
|
|
79
|
+
if len(found) != 1:
|
|
80
|
+
raise ProtocolError("missing_or_repeated_field")
|
|
81
|
+
return found[0]
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def optional(data: bytes, number: int, default=None):
|
|
85
|
+
found = [v for n, _, v, _ in fields(data) if n == number]
|
|
86
|
+
if len(found) > 1:
|
|
87
|
+
raise ProtocolError("repeated_scalar_field")
|
|
88
|
+
return found[0] if found else default
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def pb(number: int, value: bytes) -> bytes:
|
|
92
|
+
if not isinstance(value, bytes):
|
|
93
|
+
raise ProtocolError("expected_bytes")
|
|
94
|
+
return encode_varint(number * 8 + 2) + encode_varint(len(value)) + value
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def replace(data: bytes, path: list[int], value) -> bytes:
|
|
98
|
+
out = []
|
|
99
|
+
found = False
|
|
100
|
+
for number, wire, current, raw in fields(data):
|
|
101
|
+
if number != path[0]:
|
|
102
|
+
out.append(raw)
|
|
103
|
+
continue
|
|
104
|
+
if found:
|
|
105
|
+
raise ProtocolError("repeated_replace_target")
|
|
106
|
+
found = True
|
|
107
|
+
if len(path) > 1:
|
|
108
|
+
if wire != 2:
|
|
109
|
+
raise ProtocolError("expected_nested_message")
|
|
110
|
+
current = replace(current, path[1:], value)
|
|
111
|
+
else:
|
|
112
|
+
current = value
|
|
113
|
+
tag = encode_varint(number * 8 + wire)
|
|
114
|
+
if wire == 0:
|
|
115
|
+
out.append(tag + encode_varint(current))
|
|
116
|
+
elif wire == 2:
|
|
117
|
+
out.append(tag + encode_varint(len(current)) + current)
|
|
118
|
+
else:
|
|
119
|
+
raise ProtocolError("unsupported_replace_type")
|
|
120
|
+
if not found:
|
|
121
|
+
raise ProtocolError("missing_replace_target")
|
|
122
|
+
return b"".join(out)
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def frame(message: bytes) -> bytes:
|
|
126
|
+
if len(message) > MAX_FRAME:
|
|
127
|
+
raise ProtocolError("frame_too_large")
|
|
128
|
+
return b"\0" + len(message).to_bytes(4, "big") + message
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
class FrameReader:
|
|
132
|
+
def __init__(self):
|
|
133
|
+
self.buffer = bytearray()
|
|
134
|
+
self.total = 0
|
|
135
|
+
|
|
136
|
+
def feed(self, chunk: bytes) -> list[bytes]:
|
|
137
|
+
self.total += len(chunk)
|
|
138
|
+
if self.total > 32 * 1024 * 1024:
|
|
139
|
+
raise ProtocolError("response_too_large")
|
|
140
|
+
self.buffer.extend(chunk)
|
|
141
|
+
out = []
|
|
142
|
+
while len(self.buffer) >= 5:
|
|
143
|
+
size = int.from_bytes(self.buffer[1:5], "big")
|
|
144
|
+
if self.buffer[0] or size > MAX_FRAME:
|
|
145
|
+
raise ProtocolError("compressed_or_oversized_frame")
|
|
146
|
+
if len(self.buffer) < size + 5:
|
|
147
|
+
break
|
|
148
|
+
out.append(bytes(self.buffer[5:5 + size]))
|
|
149
|
+
del self.buffer[:5 + size]
|
|
150
|
+
return out
|
|
151
|
+
|
|
152
|
+
def finish(self):
|
|
153
|
+
if self.buffer:
|
|
154
|
+
raise ProtocolError("truncated_grpc_message")
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def messages(data: bytes) -> list[bytes]:
|
|
158
|
+
reader = FrameReader()
|
|
159
|
+
result = reader.feed(data)
|
|
160
|
+
reader.finish()
|
|
161
|
+
return result
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def struct_value(value: Any, depth=0) -> bytes:
|
|
165
|
+
if depth > 64:
|
|
166
|
+
raise ProtocolError("struct_too_deep")
|
|
167
|
+
if value is None:
|
|
168
|
+
return b"\x08\x00"
|
|
169
|
+
if isinstance(value, bool):
|
|
170
|
+
return b"\x20" + bytes([int(value)])
|
|
171
|
+
if isinstance(value, (int, float)):
|
|
172
|
+
if not math.isfinite(value) or (float(value).is_integer() and abs(value) > 2**53 - 1):
|
|
173
|
+
raise ProtocolError("unsafe_struct_number")
|
|
174
|
+
return b"\x11" + struct.pack("<d", value)
|
|
175
|
+
if isinstance(value, str):
|
|
176
|
+
return pb(3, value.encode())
|
|
177
|
+
if isinstance(value, dict):
|
|
178
|
+
return pb(5, struct_obj(value, depth + 1))
|
|
179
|
+
if isinstance(value, list):
|
|
180
|
+
return pb(6, b"".join(pb(1, struct_value(v, depth + 1)) for v in value))
|
|
181
|
+
raise ProtocolError("unsupported_struct_value")
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def struct_obj(value: dict, depth=0) -> bytes:
|
|
185
|
+
if depth > 64 or not all(isinstance(k, str) for k in value):
|
|
186
|
+
raise ProtocolError("invalid_struct_object")
|
|
187
|
+
return b"".join(pb(1, pb(1, k.encode()) + pb(2, struct_value(v, depth + 1)))
|
|
188
|
+
for k, v in value.items())
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
def decode_value(data: bytes, depth=0):
|
|
192
|
+
fs = fields(data)
|
|
193
|
+
if depth > 64 or len(fs) != 1:
|
|
194
|
+
raise ProtocolError("invalid_struct_value")
|
|
195
|
+
n, w, v, _ = fs[0]
|
|
196
|
+
if n == 1 and w == 0 and v == 0:
|
|
197
|
+
return None
|
|
198
|
+
if n == 2 and w == 1:
|
|
199
|
+
number = struct.unpack("<d", v)[0]
|
|
200
|
+
if not math.isfinite(number) or (number.is_integer() and abs(number) > 2**53 - 1):
|
|
201
|
+
raise ProtocolError("unsafe_struct_number")
|
|
202
|
+
return int(number) if number.is_integer() else number
|
|
203
|
+
if n == 3 and w == 2:
|
|
204
|
+
return v.decode()
|
|
205
|
+
if n == 4 and w == 0 and v in (0, 1):
|
|
206
|
+
return bool(v)
|
|
207
|
+
if n == 5 and w == 2:
|
|
208
|
+
return decode_struct(v, depth + 1)
|
|
209
|
+
if n == 6 and w == 2:
|
|
210
|
+
out = []
|
|
211
|
+
for k, wire, item, _ in fields(v):
|
|
212
|
+
if k != 1 or wire != 2:
|
|
213
|
+
raise ProtocolError("invalid_struct_list")
|
|
214
|
+
out.append(decode_value(item, depth + 1))
|
|
215
|
+
return out
|
|
216
|
+
raise ProtocolError("unknown_struct_value")
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
def decode_struct(data: bytes, depth=0) -> dict:
|
|
220
|
+
if depth > 64:
|
|
221
|
+
raise ProtocolError("struct_too_deep")
|
|
222
|
+
out = {}
|
|
223
|
+
for n, w, item, _ in fields(data):
|
|
224
|
+
if n != 1 or w != 2:
|
|
225
|
+
raise ProtocolError("invalid_struct_entry")
|
|
226
|
+
key = one(item, 1).decode()
|
|
227
|
+
if key in out:
|
|
228
|
+
raise ProtocolError("duplicate_struct_key")
|
|
229
|
+
out[key] = decode_value(one(item, 2), depth + 1)
|
|
230
|
+
return out
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
def validate_task_template(template: bytes) -> list[bytes]:
|
|
234
|
+
parts = messages(template)
|
|
235
|
+
if len(parts) != 4 or [[n for n, _, _, _ in fields(m)] for m in parts] != [[1000], [1008], [1007], [1007]]:
|
|
236
|
+
raise ProtocolError("unsupported_task_template")
|
|
237
|
+
initial = one(parts[0], 1000)
|
|
238
|
+
generate = one(one(initial, 1001), 1)
|
|
239
|
+
if (one(generate, 26) != 41 or one(generate, 49) != 1 or one(generate, 47) != 8
|
|
240
|
+
or one(one(initial, 1), 16) != MODE_ID.encode()
|
|
241
|
+
or any(one(one(generate, 4), n) for n in (1, 2, 3))):
|
|
242
|
+
raise ProtocolError("template_must_be_fresh_argon_task")
|
|
243
|
+
if one(one(generate, 1), 1) != one(one(parts[2], 1007), 1):
|
|
244
|
+
raise ProtocolError("template_prompt_copies_differ")
|
|
245
|
+
old_id = one(initial, 2)
|
|
246
|
+
if len(old_id) != 38:
|
|
247
|
+
raise ProtocolError("invalid_template_request_id")
|
|
248
|
+
uuid.UUID(old_id[2:].decode())
|
|
249
|
+
return parts
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
def build_task_request(template: bytes, prompt: str) -> bytes:
|
|
253
|
+
parts = validate_task_template(template)
|
|
254
|
+
old_id = one(one(parts[0], 1000), 2)
|
|
255
|
+
new_id = old_id[:2] + str(uuid.uuid4()).encode()
|
|
256
|
+
parts = [p.replace(old_id, new_id) for p in parts]
|
|
257
|
+
parts[0] = replace(parts[0], [1000, 1001, 1, 1, 1], prompt.encode())
|
|
258
|
+
parts[2] = replace(parts[2], [1007, 1], prompt.encode())
|
|
259
|
+
return b"".join(frame(p) for p in parts)
|
|
260
|
+
|
|
261
|
+
|
|
262
|
+
def validate_headers(headers: list, expected_path: str):
|
|
263
|
+
if not isinstance(headers, list) or not all(isinstance(h, list) and len(h) == 2
|
|
264
|
+
and all(isinstance(v, str) and '\r' not in v and '\n' not in v for v in h) for h in headers):
|
|
265
|
+
raise ProtocolError("invalid_headers")
|
|
266
|
+
h = dict(headers)
|
|
267
|
+
if len(h) != len(headers) or (h.get(":authority"), h.get(":scheme"), h.get(":method"), h.get(":path"),
|
|
268
|
+
h.get("content-type")) != (HOST + ":443", "https", "POST", expected_path, "application/grpc"):
|
|
269
|
+
raise ProtocolError("unexpected_rpc_target")
|
|
270
|
+
if expected_path == TASK_RPC and not h.get("x-server-token"):
|
|
271
|
+
raise ProtocolError("missing_server_token")
|
|
272
|
+
|
|
273
|
+
|
|
274
|
+
def turn_request(template: bytes, conversation: bytes) -> bytes:
|
|
275
|
+
parts = messages(template)
|
|
276
|
+
if len(parts) != 1 or not one(parts[0], 1).startswith(b"c_") or one(parts[0], 2) != 100:
|
|
277
|
+
raise ProtocolError("unsupported_turn_template")
|
|
278
|
+
if not conversation.startswith(b"c_") or len(conversation) > 100:
|
|
279
|
+
raise ProtocolError("invalid_conversation")
|
|
280
|
+
return frame(replace(replace(parts[0], [1], conversation), [2], 1))
|
|
281
|
+
|
|
282
|
+
|
|
283
|
+
def tool_reply(function: bytes, context: bytes, payload: dict) -> bytes:
|
|
284
|
+
return pb(2, pb(1, pb(1, function) + pb(2, context) + pb(3, struct_obj(payload))))
|
|
285
|
+
|
|
286
|
+
|
|
287
|
+
@dataclass
|
|
288
|
+
class TaskOutput:
|
|
289
|
+
conversation: bytes | None = None
|
|
290
|
+
request: bytes | None = None
|
|
291
|
+
text: str = ""
|
|
292
|
+
completeness: int | None = None
|
|
293
|
+
mode_id: str | None = None
|
|
294
|
+
attribution: str | None = None
|
|
295
|
+
downgrade: int = 0
|
|
296
|
+
|
|
297
|
+
def feed(self, message: bytes):
|
|
298
|
+
for n, w, value, _ in fields(message):
|
|
299
|
+
if n == 1014 and w == 2:
|
|
300
|
+
candidate = one(one(value, 1), 1)
|
|
301
|
+
if not candidate.startswith(b"c_") or len(candidate) > 100:
|
|
302
|
+
raise ProtocolError("invalid_conversation")
|
|
303
|
+
if self.conversation is not None and self.conversation != candidate:
|
|
304
|
+
raise ProtocolError("conversation_changed")
|
|
305
|
+
self.conversation = candidate
|
|
306
|
+
if n != 1003 or w != 2:
|
|
307
|
+
continue
|
|
308
|
+
for kind, wire, response, _ in fields(value):
|
|
309
|
+
if kind != 1 or wire != 2:
|
|
310
|
+
continue
|
|
311
|
+
ids = optional(response, 2)
|
|
312
|
+
if ids is not None:
|
|
313
|
+
c, r = one(ids, 1), one(ids, 2)
|
|
314
|
+
if c and r:
|
|
315
|
+
if self.conversation not in (None, c) or self.request not in (None, r):
|
|
316
|
+
raise ProtocolError("response_identifiers_changed")
|
|
317
|
+
self.conversation, self.request = c, r
|
|
318
|
+
for field, attr in [(25, "mode_id"), (30, "attribution"), (31, "downgrade")]:
|
|
319
|
+
v = optional(response, field)
|
|
320
|
+
if v is not None:
|
|
321
|
+
setattr(self, attr, v.decode() if isinstance(v, bytes) else v)
|
|
322
|
+
model = optional(response, 1)
|
|
323
|
+
if model is not None:
|
|
324
|
+
self._model(model)
|
|
325
|
+
|
|
326
|
+
def _model(self, model: bytes):
|
|
327
|
+
# Never combine an earlier partial text with a later empty terminal
|
|
328
|
+
# message. A usable terminal snapshot must supply both in one model.
|
|
329
|
+
self.completeness = optional(model, 19)
|
|
330
|
+
self.text = ""
|
|
331
|
+
display = optional(model, 2)
|
|
332
|
+
if display is not None:
|
|
333
|
+
text = optional(display, 1)
|
|
334
|
+
if text is not None:
|
|
335
|
+
self.text = text.decode()
|
|
336
|
+
|
|
337
|
+
def feed_saved(self, data: bytes) -> bool:
|
|
338
|
+
if not self.conversation or not self.request:
|
|
339
|
+
raise ProtocolError("missing_response_identifiers")
|
|
340
|
+
found = False
|
|
341
|
+
for message in messages(data):
|
|
342
|
+
for n, w, turn, _ in fields(message):
|
|
343
|
+
if n != 1 or w != 2:
|
|
344
|
+
continue
|
|
345
|
+
ids = one(turn, 1)
|
|
346
|
+
if (one(ids, 1), one(ids, 2)) != (self.conversation, self.request):
|
|
347
|
+
continue
|
|
348
|
+
if found:
|
|
349
|
+
raise ProtocolError("duplicate_saved_turn")
|
|
350
|
+
found = True
|
|
351
|
+
assistant = one(turn, 4)
|
|
352
|
+
self._model(one(assistant, 1))
|
|
353
|
+
return found
|
|
354
|
+
|
|
355
|
+
def check_route(self, final=False):
|
|
356
|
+
if self.downgrade or self.mode_id not in (None, MODE_ID):
|
|
357
|
+
raise ProtocolError("upstream_model_changed")
|
|
358
|
+
if final and self.mode_id != MODE_ID:
|
|
359
|
+
raise ProtocolError("upstream_model_unverified")
|
|
360
|
+
|
|
361
|
+
@property
|
|
362
|
+
def complete(self) -> bool:
|
|
363
|
+
return self.completeness == 2 and bool(self.text.strip())
|