|
| 1 | +package lip |
| 2 | + |
| 3 | +import ( |
| 4 | + "encoding/json" |
| 5 | + "io" |
| 6 | + "net" |
| 7 | + "time" |
| 8 | +) |
| 9 | + |
| 10 | +// StreamContextSymbol is one frame of a StreamContext response. The |
| 11 | +// embedded `OwnedSymbolInfo` is flattened into the fields we actually |
| 12 | +// consume in CKB — the full Rust struct carries many fields we don't |
| 13 | +// need (telemetry, relationships, taint) and serialising them through |
| 14 | +// `map[string]any` would be wasteful. |
| 15 | +type StreamContextSymbol struct { |
| 16 | + URI string `json:"uri"` |
| 17 | + DisplayName string `json:"display_name"` |
| 18 | + Kind string `json:"kind"` |
| 19 | + RelevanceScore float32 `json:"relevance_score"` |
| 20 | + TokenCost uint32 `json:"token_cost"` |
| 21 | +} |
| 22 | + |
| 23 | +// StreamContextResult summarises a completed StreamContext stream. |
| 24 | +// `Reason` is one of "token_budget" | "exhausted" | "error". |
| 25 | +type StreamContextResult struct { |
| 26 | + Symbols []StreamContextSymbol |
| 27 | + Reason string |
| 28 | + Emitted uint32 |
| 29 | + TotalCandidates uint32 |
| 30 | + Err string |
| 31 | +} |
| 32 | + |
| 33 | +// StreamContextPosition is the cursor rectangle the daemon ranks around. |
| 34 | +// Byte-offset semantics match LIP's `OwnedRange` — 0-based lines and |
| 35 | +// chars. Pass a zero-width range at the cursor, or a whole-file range |
| 36 | +// (`start_line=0, end_line=lineCount`) for file-level context. |
| 37 | +type StreamContextPosition struct { |
| 38 | + StartLine int `json:"start_line"` |
| 39 | + StartChar int `json:"start_char"` |
| 40 | + EndLine int `json:"end_line"` |
| 41 | + EndChar int `json:"end_char"` |
| 42 | +} |
| 43 | + |
| 44 | +// streamContextMaxFrames caps how many SymbolInfo frames we accept before |
| 45 | +// bailing out — defence against a runaway daemon. Large indexes could |
| 46 | +// theoretically produce 10k+ candidates; a hard cap of 1024 is well above |
| 47 | +// any realistic caller budget and cheap to enforce. |
| 48 | +const streamContextMaxFrames = 1024 |
| 49 | + |
| 50 | +// StreamContext opens a dedicated connection, sends a `stream_context` |
| 51 | +// request, and reads SymbolInfo frames until the daemon writes the |
| 52 | +// `end_stream` terminator. Returns (nil, nil) when the daemon is |
| 53 | +// unavailable — callers must treat nil as "LIP unavailable" (same contract |
| 54 | +// as the rest of the package). |
| 55 | +// |
| 56 | +// The dedicated connection is intentional: `stream_context` on the |
| 57 | +// shared subscriber channel would interleave with IndexStatus pings and |
| 58 | +// IndexChanged pushes and complicate parsing. One connection per call is |
| 59 | +// fine — the ranking itself dominates latency, and callers shouldn't |
| 60 | +// issue this RPC more than a few times per second. |
| 61 | +func StreamContext(fileURI string, pos StreamContextPosition, maxTokens uint32, model string) (*StreamContextResult, error) { |
| 62 | + conn, err := net.DialTimeout("unix", SocketPath(), 500*time.Millisecond) |
| 63 | + if err != nil { |
| 64 | + return nil, nil |
| 65 | + } |
| 66 | + defer conn.Close() |
| 67 | + // Overall deadline: the daemon's relevance ranking is heuristic and |
| 68 | + // bounded, but pathological inputs could stall. 10 s is generous; for |
| 69 | + // a token_budget of ~2000 it completes in ~200 ms typically. |
| 70 | + _ = conn.SetDeadline(time.Now().Add(10 * time.Second)) |
| 71 | + |
| 72 | + req := map[string]any{ |
| 73 | + "type": "stream_context", |
| 74 | + "file_uri": fileURI, |
| 75 | + "cursor_position": pos, |
| 76 | + "max_tokens": maxTokens, |
| 77 | + } |
| 78 | + if model != "" { |
| 79 | + req["model"] = model |
| 80 | + } |
| 81 | + if err := writeFrame(conn, req); err != nil { |
| 82 | + return nil, nil |
| 83 | + } |
| 84 | + |
| 85 | + out := &StreamContextResult{Symbols: make([]StreamContextSymbol, 0, 16)} |
| 86 | + for range streamContextMaxFrames + 1 { |
| 87 | + frame, err := readFrame(conn) |
| 88 | + if err != nil { |
| 89 | + if err == io.EOF { |
| 90 | + return out, nil |
| 91 | + } |
| 92 | + return nil, nil |
| 93 | + } |
| 94 | + // ServerResponse { ok: ServerMessage, error: Option<String> } |
| 95 | + inner := frame |
| 96 | + if raw, ok := frame["ok"]; ok && len(raw) > 0 && string(raw) != "null" { |
| 97 | + _ = json.Unmarshal(raw, &inner) |
| 98 | + } |
| 99 | + var kind string |
| 100 | + _ = json.Unmarshal(inner["type"], &kind) |
| 101 | + |
| 102 | + switch kind { |
| 103 | + case "symbol_info": |
| 104 | + var sym struct { |
| 105 | + SymbolInfo struct { |
| 106 | + URI string `json:"uri"` |
| 107 | + DisplayName string `json:"display_name"` |
| 108 | + Kind string `json:"kind"` |
| 109 | + } `json:"symbol_info"` |
| 110 | + RelevanceScore float32 `json:"relevance_score"` |
| 111 | + TokenCost uint32 `json:"token_cost"` |
| 112 | + } |
| 113 | + if b, ok := marshalInner(inner); ok { |
| 114 | + _ = json.Unmarshal(b, &sym) |
| 115 | + } |
| 116 | + out.Symbols = append(out.Symbols, StreamContextSymbol{ |
| 117 | + URI: sym.SymbolInfo.URI, |
| 118 | + DisplayName: sym.SymbolInfo.DisplayName, |
| 119 | + Kind: sym.SymbolInfo.Kind, |
| 120 | + RelevanceScore: sym.RelevanceScore, |
| 121 | + TokenCost: sym.TokenCost, |
| 122 | + }) |
| 123 | + case "end_stream": |
| 124 | + var end struct { |
| 125 | + Reason string `json:"reason"` |
| 126 | + Emitted uint32 `json:"emitted"` |
| 127 | + TotalCandidates uint32 `json:"total_candidates"` |
| 128 | + Error *string `json:"error"` |
| 129 | + } |
| 130 | + if b, ok := marshalInner(inner); ok { |
| 131 | + _ = json.Unmarshal(b, &end) |
| 132 | + } |
| 133 | + out.Reason = end.Reason |
| 134 | + out.Emitted = end.Emitted |
| 135 | + out.TotalCandidates = end.TotalCandidates |
| 136 | + if end.Error != nil { |
| 137 | + out.Err = *end.Error |
| 138 | + } |
| 139 | + return out, nil |
| 140 | + case "error", "unknown_message": |
| 141 | + // Daemon rejected the request — treat as unavailable. |
| 142 | + return nil, nil |
| 143 | + default: |
| 144 | + // Unknown frame type mid-stream: skip rather than fail hard. |
| 145 | + } |
| 146 | + } |
| 147 | + return out, nil |
| 148 | +} |
0 commit comments